diff --git a/tests/test_server/test_grpc/test_fligh_control_servicer.py b/tests/test_server/test_grpc/test_fligh_control_servicer.py
index 52f32c41..363987e2 100644
--- a/tests/test_server/test_grpc/test_fligh_control_servicer.py
+++ b/tests/test_server/test_grpc/test_fligh_control_servicer.py
@@ -29,10 +29,10 @@ def _pinger():
     pinger_thread.start()
 
     assert not evt.is_set()
-    assert not evt.wait(timeout=0.2)
+    assert not evt.wait(timeout=1)
 
     stop_pinger.set()
-    assert evt.wait(timeout=0.2)
+    assert evt.wait(timeout=1)
 
 
 def test_shutdown_timeout_0_means_no_watchdog():
diff --git a/tests/test_server/test_grpc/test_training_servicer.py b/tests/test_server/test_grpc/test_training_servicer.py
index b40d367e..5de85d2a 100644
--- a/tests/test_server/test_grpc/test_training_servicer.py
+++ b/tests/test_server/test_grpc/test_training_servicer.py
@@ -107,6 +107,7 @@ def unet2d_config_path(checkpoint_dir, train_data_dir, val_data_path, device: st
           execution_probability: 1.0
         - name: ToTensor
           expand_dims: true
+          dtype: 'long'
       label:
         - name: RandomFlip
         - name: RandomRotate90
@@ -125,6 +126,7 @@ def unet2d_config_path(checkpoint_dir, train_data_dir, val_data_path, device: st
       weight:
         - name: ToTensor
           expand_dims: false
+          dtype: 'long'
   val:
     file_paths:
       - {val_data_path}
@@ -140,6 +142,7 @@ def unet2d_config_path(checkpoint_dir, train_data_dir, val_data_path, device: st
         - name: Standardize
         - name: ToTensor
           expand_dims: true
+          dtype: 'long'
       label:
         - name: ToTensor
           expand_dims: false
@@ -147,6 +150,7 @@ def unet2d_config_path(checkpoint_dir, train_data_dir, val_data_path, device: st
       weight:
         - name: ToTensor
           expand_dims: false
+          dtype: 'long'
 """