Update train_loop.py

Author: mainpyp

src/train_infer/train_loop.py

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