test: add test for state_dict and b

jmclong · Oct 22, 2022 · 836960b · 836960b
1 parent ea9b07d
commit 836960b
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/test/test_layers.py b/test/test_layers.py
@@ -29,6 +29,18 @@ def test_gaussian_encoding(device):
         decimal=5)
 
 
+@pytest.mark.parametrize('device', ['cpu', 'cuda'])
+def test_gaussian_encoding_register_buffer(device):
+    check_cuda(device)
+    b = rff.functional.sample_b(1.0, (256, 2)).to(device)
+    layer = rff.layers.GaussianEncoding(b=b).to(device)
+    assert 'b' in layer.state_dict()
+    np.testing.assert_almost_equal(
+        layer.state_dict()['b'].cpu().numpy(),
+        b.cpu().numpy(),
+        decimal=5)
+
+
 @pytest.mark.parametrize('device', ['cpu', 'cuda'])
 def test_basic_encoding(device):
     check_cuda(device)