Update main.py

yjyjy131 · Jul 7, 2022 · ae4d2f5 · ae4d2f5
1 parent a239cf6
commit ae4d2f5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/main.py b/main.py
@@ -11,11 +11,11 @@
 
     ### ORIGINAL VERSION ###
     # Train model
-    slot_network_training(data_num=10, batch_size=12, valid_data_num=1500, valid_batch_size=32, epoch=10, input_res=224, device_id=0, num_workers=0)
+    slot_network_training(data_num=10, batch_size=12, valid_data_num=1500, valid_batch_size=32, epoch=10, input_res=224, device_id=0, num_workers=1)
 
     # Test model 
     params_path = './parameters/merge_bn_epoch10_loss4.pkl'
-    slot_network_testing(parameter_path=params_path, data_num=1500, batch_size=50, input_res=224, device_id=0,  num_workers=0)
+    slot_network_testing(parameter_path=params_path, data_num=1500, batch_size=50, input_res=224, device_id=0,  num_workers=1)
 
     # Load detector
     detector = SlotDetector(device_id=0, dim_encoder=[32, 44, 64, 92, 128], parameter_path=params_path)