mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-08-13 21:55:46 +00:00
Feature/vit support (#4182)
* [shardformer] added tests * [shardformer] vit test finish and support * fix attention dropout
This commit is contained in:
parent
0ceec8f9a9
commit
c59d7aca09
@ -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
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user