add Exception('Batchsize must be greater than num_gpus.')

Author: zylo117

utils/utils.py

@@ -195,6 +195,9 @@ def scatter(self, inputs, kwargs, device_ids):
         devices = ['cuda:' + str(x) for x in range(self.num_gpus)]
         splits = inputs[0].shape[0] // self.num_gpus
 
+        if splits == 0:
+            raise Exception('Batchsize must be greater than num_gpus.')
+
         return [(inputs[0][splits * device_idx: splits * (device_idx + 1)].to(f'cuda:{device_idx}', non_blocking=True),
                  inputs[1][splits * device_idx: splits * (device_idx + 1)].to(f'cuda:{device_idx}', non_blocking=True))
                 for device_idx in range(len(devices))], \