From 06422fbe19dbcdaac6d6fb0983bcc980a5568ea9 Mon Sep 17 00:00:00 2001 From: "Deng, Weishi" Date: Tue, 12 Dec 2023 19:47:36 -0800 Subject: [PATCH] add xpu profile info --- torchbenchmark/models/nvidia_deeprecommender/nvinfer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/torchbenchmark/models/nvidia_deeprecommender/nvinfer.py b/torchbenchmark/models/nvidia_deeprecommender/nvinfer.py index 34116270a3..90418b7ac6 100644 --- 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: