github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-13 21:55:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feature/vit support (#4182)

Browse Source 
* [shardformer] added tests

* [shardformer] vit test finish and support

* fix attention dropout
This commit is contained in:
Kun Lin 2023-07-07 14:06:46 +08:00 committed by Hongxin Liu
parent 0ceec8f9a9
commit c59d7aca09
1 changed files with 42 additions and 41 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/shardformer/policies/vit.py
View File

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