amd
/

Nitro-T-0.6B

akasharidas commited on Jun 25

Commit

6cde8f0

verified ·

1 Parent(s): 8c8069e

Update transformer/transformer.py

Files changed (1) hide show

transformer/transformer.py CHANGED Viewed

@@ -100,7 +100,7 @@ class RMSNorm(torch.nn.Module):
 # Modified from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/pixart_transformer_2d.py
-class EspressoDiTModel(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
     _no_split_modules = ["BasicTransformerBlock", "PatchEmbed"]

 # Modified from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/pixart_transformer_2d.py
+class NitroDiTModel(ModelMixin, ConfigMixin):
     _supports_gradient_checkpointing = True
     _no_split_modules = ["BasicTransformerBlock", "PatchEmbed"]