Update modeling_tiny_mixtral.py
Browse files- modeling_tiny_mixtral.py +1 -0
modeling_tiny_mixtral.py
CHANGED
|
@@ -582,6 +582,7 @@ class TinyMixtralForCausalLM(PreTrainedModel, GenerationMixin):
|
|
| 582 |
)
|
| 583 |
self.model = tiny_mixtral(args=args)
|
| 584 |
self.config = config
|
|
|
|
| 585 |
self.post_init()
|
| 586 |
|
| 587 |
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
|
|
|
|
| 582 |
)
|
| 583 |
self.model = tiny_mixtral(args=args)
|
| 584 |
self.config = config
|
| 585 |
+
self.model.to(config.device)
|
| 586 |
self.post_init()
|
| 587 |
|
| 588 |
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
|