Merge pull request #9 from kengz/1.1.2

set encoding as nn.Parameter
kengz · Dec 31, 2021 · cac3dbd · cac3dbd
2 parents 1f60d2f + fcbf0da
commit cac3dbd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -32,7 +32,7 @@ def run_tests(self):
 
 setup(
  name='torcharc',
- version='1.1.1',
+ version='1.1.2',
  description='Build PyTorch networks by specifying architectures.',
  long_description='https://github.com/kengz/torcharc',
  keywords='torcharc',

diff --git a/torcharc/module/perceiver_io/preprocessor.py b/torcharc/module/perceiver_io/preprocessor.py
@@ -95,7 +95,7 @@ def build_pos_encoding(self, pos: torch.Tensor, max_reso: list = None) -> torch.
  spatial_encoding = torch.cat(encodings, dim=-1) # shape (x, y,... d*(2*num_freq_bands+1))
  # flatten spatial dimensions into 1D
  pos_encoding = rearrange(spatial_encoding, '... c -> (...) c')
- return pos_encoding
+ return nn.Parameter(pos_encoding)
 
  def get_pos_encoding_dim(self) -> int:
  return len(self.spatial_shape) * (2 * self.num_freq_bands + int(self.cat_pos))