move model to cuda

pytorch · Dec 18, 2024 · aac32ae · aac32ae
1 parent a239d0d
commit aac32ae
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tools/perf/perf_run.py b/tools/perf/perf_run.py
@@ -243,7 +243,6 @@ def run_dynamo(model, input_tensors, params, precision, batch_size):
         " batch_size : ",
         batch_size,
     )
-    model = model.to("cuda:0")
     if params["is_text_llm"]:
         return run_hf_dynamo(model, input_tensors, params, precision, batch_size)
 
@@ -638,9 +637,9 @@ def run(
     # Load PyTorch Model, if provided
     if len(model_name_torch) > 0 and os.path.exists(model_name_torch):
         print("Loading user provided torch model: ", model_name_torch)
-        model_torch = torch.load(model_name_torch).eval()
+        model_torch = torch.load(model_name_torch).cuda().eval()
     elif model_name_torch in BENCHMARK_MODELS:
-        model_torch = BENCHMARK_MODELS[model_name_torch]["model"].eval()
+        model_torch = BENCHMARK_MODELS[model_name_torch]["model"].cuda().eval()
 
     # If neither model type was provided
     if (model is None) and (model_torch is None):