debug: show cuda devices

jolibrain · Nov 27, 2024 · 8f3fd76 · 8f3fd76
1 parent 0e53038
commit 8f3fd76
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/scripts/gen_single_image.py b/scripts/gen_single_image.py
@@ -56,6 +56,12 @@ def load_model(modelpath, model_in_file, cpu, gpuid):
     else:
         model.load_state_dict(torch.load(modelpath + "/" + model_in_file))
 
+    print("cuda is_available=", torch.cuda.is_available())
+    print("device_count=", torch.cuda.device_count())
+    print("device=", torch.cuda.current_device())
+    print("cuda device name=", torch.cuda.get_device_name(0))
+    print("cuda device=", torch.cuda.device(0))
+
     model = model.to(device)
     return model, opt, device