File size: 982 Bytes
07fcdfe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | import unittest
from pdgrapher import Trainer
import torch
import os
torch.set_num_threads(5)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
class TestTrainer(unittest.TestCase):
def test_trainer(self):
trainer = Trainer()
def test_trainer_args(self):
trainer = Trainer(
fabric_kwargs={"accelerator": "cpu"},
log_train=False, use_backward_data=True, supervision_multiplier=1.3
)
self.assertFalse(trainer.log_train)
self.assertTrue(trainer.use_backward_data)
self.assertEqual(trainer.supervision_multiplier, 1.3)
self.assertEqual(trainer.fabric.device.type, "cpu")
def test_trainer_unknown_kwargs(self):
with self.assertWarns(Warning):
trainer = Trainer(
unknown_arg=1, another_unknown=True
)
if __name__ == "__main__":
unittest.TestLoader.sortTestMethodsUsing = None
unittest.main()
|