fix small bugs

pytorch · Apr 16, 2024 · 2d40091 · 2d40091
1 parent 645513e
commit 2d40091
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/embedding.py b/py/torch_tensorrt/dynamo/conversion/impl/embedding.py
@@ -360,7 +360,7 @@ def embedding_bag(
         weight,
     )
     embed = cast_trt_tensor(
-        ctx, embed, torch.float, f"{name}_cast_embed_to_fp16", target, source_ir
+        ctx, embed, torch.float, f"{name}_cast_embed_to_fp32", target, source_ir
     )
 
     # give weights to embedding

diff --git a/tests/py/dynamo/conversion/harness.py b/tests/py/dynamo/conversion/harness.py
@@ -92,14 +92,14 @@ def run_test(
             ):
                 ref_outputs = [ref_outputs]
             for out, ref in zip(outputs, ref_outputs):
+                ref = ref.cpu()  # to_dtype test has cases with gpu output
                 if not isinstance(ref, torch.Tensor):
                     ref = torch.tensor([ref])
-                ref = ref.cpu()  # to_dtype test has cases with gpu output
                 if ref.dtype == torch.int64:
                     ref = ref.int()  # convert torch.max's index output tensor to int32
                 torch.testing.assert_close(
                     out.cpu(),
-                    ref.cpu(),
+                    ref,
                     rtol=rtol,
                     atol=atol,
                     equal_nan=True,