diff --git a/tests/test_layers.py b/tests/test_layers.py
index <HASH>..<HASH> 100644
--- a/tests/test_layers.py
+++ b/tests/test_layers.py
@@ -240,7 +240,15 @@ def test_torch_layer():
         assert not torch.allclose(model1(input), model2(input))
         model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict())))
         assert torch.allclose(model1(input), model2(input))
-        # TODO jittization
+
+        # tracing (freezing)
+        model3 = torch.jit.trace(input)(model2)
+        assert torch.allclose(model1(input), model3(input))
+        assert torch.allclose(model1(input + 1), model3(input + 1))
+
+        model4 = torch.jit.trace(input, optimize=True)(model2)
+        assert torch.allclose(model1(input), model4(input))
+        assert torch.allclose(model1(input + 1), model4(input + 1))
 
 
 def test_keras_layer():
Add testing of model freezing for pytorch