Wrap outs.to_dense() inside if layout=_mkldnn

undertherain · May 11, 2021 · 05d1611 · 05d1611
1 parent c012795
commit 05d1611
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/benchmarker/kernels/helpers_torch.py b/benchmarker/kernels/helpers_torch.py
@@ -27,7 +27,8 @@ def __call__(self, x, t):
         # it?
         if isinstance(outs, OrderedDict):
             outs = outs["out"]
-        outs = outs.to_dense()
+        if outs.layout == torch._mkldnn:
+            outs = outs.to_dense()
         loss = self.criterion(outs, t)
         return loss