Feature/vit support (#4182)

* [shardformer] added tests

* [shardformer] vit test finish and support

* fix attention dropout

colossalai/shardformer/policies/vit.py

@@ -1,4 +1,3 @@
-from functools import partial
 from typing import Callable, Dict, List, Union
 
 import torch.nn as nn
@@ -86,6 +85,8 @@ class ViTPolicy(Policy):
 
         return policy
 
+        return policy
+
     def new_model_class(self):
         return None