| from torch import nn | |
| FC_CLASS_REGISTRY = {'torch': nn.Linear} | |
| try: | |
| import transformer_engine.pytorch as te | |
| FC_CLASS_REGISTRY['te'] = te.Linear | |
| except: | |
| pass |
| from torch import nn | |
| FC_CLASS_REGISTRY = {'torch': nn.Linear} | |
| try: | |
| import transformer_engine.pytorch as te | |
| FC_CLASS_REGISTRY['te'] = te.Linear | |
| except: | |
| pass |