add xpu profile info

pytorch · Dec 13, 2023 · 06422fb · 06422fb
1 parent b5620ff
commit 06422fb
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/torchbenchmark/models/nvidia_deeprecommender/nvinfer.py b/torchbenchmark/models/nvidia_deeprecommender/nvinfer.py
@@ -136,7 +136,6 @@ def __init__(self, device = 'cpu', jit=False, batch_size=256, usecommandlineargs
     if self.toytest:
       self.toyinputs = torch.randn(self.batch_size,self.node_count).to(device)
 
-    print("device:", device)
     if usecommandlineargs:
       self.args = getCommandLineArgs()
     else:
@@ -253,7 +252,7 @@ def TimedInferenceRun(self) :
       e_start_time = time.time()
 
       if self.args.profile:
-        with profiler.profile(record_shapes=True, use_cuda=True) as prof:
+        with profiler.profile(record_shapes=True, use_cuda=self.args.use_cuda, use_xpu=self.args.use_xpu) as prof:
           with profiler.record_function("Inference"):
             self.eval()
       else: