Update modeling.py
Browse files- modeling.py +2 -1
modeling.py
CHANGED
|
@@ -148,4 +148,5 @@ class CustomModel(PreTrainedModel):
|
|
| 148 |
self.model = GPT(config.d_model, config.n_heads, config.d_head, config.n_vocab, config.n_layers)
|
| 149 |
|
| 150 |
def forward(self, tensor):
|
| 151 |
-
|
|
|
|
|
|
| 148 |
self.model = GPT(config.d_model, config.n_heads, config.d_head, config.n_vocab, config.n_layers)
|
| 149 |
|
| 150 |
def forward(self, tensor):
|
| 151 |
+
with torch.autocast('cuda', dtype=torch.bfloat16):
|
| 152 |
+
return self.model(tensor)
|