plt update torch, test=model

Zeref996 · Jan 2, 2025 · 3768d4b · 3768d4b
1 parent 6882db6
commit 3768d4b
Show file tree

Hide file tree

Showing 9 changed files with 468 additions and 7 deletions.
diff --git a/framework/e2e/PaddleLT_new/engine/torch_eval.py b/framework/e2e/PaddleLT_new/engine/torch_eval.py
@@ -30,7 +30,11 @@ def __init__(self, testing, layerfile, device_place_id, upstream_net, orderdict_
         reset(self.seed)
 
         self.device = os.environ.get("PLT_SET_DEVICE")
-        torch.device(f"cuda:{device_place_id}")
+        # torch.device(f"cuda:{device_place_id}")
+        device = torch.device(f"cuda:{device_place_id}")
+        # device = torch.device('cuda:0')
+        # torch.cuda.set_device(device)
+        torch.set_default_device(device)
 
         self.testing = testing
         self.upstream_net = upstream_net