From 266b470b832f22afa234ee69de8d55abeaa1c7dc Mon Sep 17 00:00:00 2001 From: Matt Tancik Date: Thu, 2 Jun 2022 09:44:50 -0700 Subject: [PATCH] Save RFF state to buffer --- pyrad/nerf/field_modules/encoding.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyrad/nerf/field_modules/encoding.py b/pyrad/nerf/field_modules/encoding.py index 1632030ea..3b65039b1 100644 --- a/pyrad/nerf/field_modules/encoding.py +++ b/pyrad/nerf/field_modules/encoding.py @@ -139,6 +139,7 @@ class RFFEncoding(Encoding): raise ValueError("RFF encoding scale should be greater than zero") self.scale = scale self.b_matrix = torch.normal(mean=0, std=self.scale, size=(self.in_dim, self.num_frequencies)) + self.register_buffer(name="rff_scalings", tensor=self.b_matrix) def get_out_dim(self) -> int: return self.num_frequencies * 2