Add conversion test

NVIDIA-Merlin · Jul 1, 2023 · ab0f635 · ab0f635
1 parent ddcf3bc
commit ab0f635
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/tests/unit/torch/test_block.py b/tests/unit/torch/test_block.py
@@ -322,3 +322,14 @@ def test_nesting(self):
 
         assert torch.equal(output["shortcut"], inputs)
         assert torch.equal(output["output"], inputs + 1)
+
+    def test_convert(self):
+        block = Block(PlusOne())
+        shortcut = ShortcutBlock(*block)
+
+        assert isinstance(shortcut[0], PlusOne)
+        inputs = torch.rand(5, 5)
+        assert torch.equal(
+            module_utils.module_test(shortcut, inputs)["output"],
+            module_utils.module_test(ShortcutBlock(PlusOne()), inputs)["output"],
+        )