Reverts device settings

Author: mainpyp

src/train_infer/train_loop.py

@@ -108,10 +108,8 @@ def __init__(
             self.use_ddp = True
             self.ddp_model = DDP(
                 self.model,
-                device_ids=[2],
-                output_device=2,
-                # device_ids=[dist_util.dev()],
-                # output_device=dist_util.dev(),
+                device_ids=[dist_util.dev()],
+                output_device=dist_util.dev(),
                 broadcast_buffers=False,
                 bucket_cap_mb=128,
                 find_unused_parameters=False,