small bug with init

lucidrains · Sep 18, 2021 · 83f64c4 · 83f64c4
1 parent 7c79998
commit 83f64c4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/se3_transformer_pytorch/se3_transformer_pytorch.py b/se3_transformer_pytorch/se3_transformer_pytorch.py
@@ -298,7 +298,7 @@ def __init__(
         self.apply(self.init_)
 
     def init_(self, m):
-        if m in {nn.Linear}:
+        if isinstance(m, nn.Linear):
             nn.init.kaiming_uniform_(m.weight)
 
     def forward(self, x):

diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
   name = 'se3-transformer-pytorch',
   packages = find_packages(),
   include_package_data = True,
-  version = '0.8.10',
+  version = '0.8.11',
   license='MIT',
   description = 'SE3 Transformer - Pytorch',
   author = 'Phil Wang',