Merge branch 'fix-batch-data-test' into 'main'

Fix batch data test

See merge request heka/medkit!239

changelog: Fix batch data test

tests/unit/training/test_batch_data.py

@@ -36,10 +36,8 @@ def test_dict_properties():
 @pytest.mark.skipif(not TEST_CUDA, reason="cuda is not available")
 def test_to_device():
     cpu = torch.device("cpu")
-    gpu = torch.device("gpu")
-    data = BatchData(
-        inputs=["hello", "world"], outputs=[torch.tensor(0), torch.tensor(1)]
-    )
+    gpu = torch.device("cuda:0")
+    data = BatchData(inputs=["hello", "world"], outputs=torch.tensor([0, 1]))
     new_data = data.to_device(gpu)
     for tensor_cpu, tensor_gpu in zip(data["outputs"], new_data["outputs"]):
         assert tensor_cpu.device == cpu