fix: test_cuda in test_trainer.py::TestTrainer
#572
Merged
test_trainer.py::TestTrainer
#572