Skip to content

GitLab

Unverified Commit 8cbd0bd1 authored by Lysandre Debut's avatar Lysandre Debut Committed by GitHub
Browse files

Specify dataset dtype (#10195) 


Co-authored-by: default avatarQuentin Lhoest <lhoest.q@gmail.com>
Co-authored-by: default avatarQuentin Lhoest <lhoest.q@gmail.com>
parent 0b1f552a
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
tests/test_trainer.py
View file @ 8cbd0bd1
...@@ -500,7 +500,7 @@ class TrainerIntegrationTest(unittest.TestCase): ...@@ -500,7 +500,7 @@ class TrainerIntegrationTest(unittest.TestCase):
self.check_trained_model(trainer.model) self.check_trained_model(trainer.model)
# Can return tensors. # Can return tensors.
train_dataset.set_format(type="torch") train_dataset.set_format(type="torch", dtype=torch.float32)
model = RegressionModel() model = RegressionModel()
trainer = Trainer(model, args, train_dataset=train_dataset) trainer = Trainer(model, args, train_dataset=train_dataset)
trainer.train() trainer.train()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment