Fix issue while logging cuda device utilization to tensorboard (#3727)

ludwig-ai · Oct 13, 2023 · 062f958 · 062f958
1 parent cb1f7d9
commit 062f958
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ludwig/trainers/trainer.py b/ludwig/trainers/trainer.py
@@ -492,7 +492,7 @@ def write_step_summary(cls, train_summary_writer, combined_loss, all_losses, ste
                 # https://pytorch.org/docs/stable/generated/torch.cuda.utilization.html#torch.cuda.utilization
                 train_summary_writer.add_scalar(
                     f"cuda/device{i}/utilization",
-                    torch.cuda.device(i).utilization(),
+                    torch.cuda.utilization(device=device),
                     global_step=step,
                 )
         train_summary_writer.flush()