fix STN init bug

kentechx · kentechx · commit a7f9ad0267c5 · 2024-09-16T15:20:16.000+08:00
diff --git a/pointnet/pointnet.py b/pointnet/pointnet.py
@@ -47,7 +47,7 @@ def __init__(self, in_dim=3, out_nd=None, head_norm=True):
         )
 
         nn.init.normal_(self.head[-1].weight, 0, 0.001)
-        nn.init.eye_(self.head[-1].bias.view(in_dim, in_dim))
+        nn.init.eye_(self.head[-1].bias.view(self.out_nd, self.out_nd))
 
     def forward(self, x):
         # x: (b, d, n)
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name='pointnet',
     packages=find_packages(),
-    version='0.1.1',
+    version='0.1.2',
     license='MIT',
     description='PointNet - Pytorch',
     author='Kaidi Shen',

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ def __init__(self, in_dim=3, out_nd=None, head_norm=True):`
`47`	`47`	`)`
`48`	`48`
`49`	`49`	`nn.init.normal_(self.head[-1].weight, 0, 0.001)`
`50`		`- nn.init.eye_(self.head[-1].bias.view(in_dim, in_dim))`
	`50`	`+ nn.init.eye_(self.head[-1].bias.view(self.out_nd, self.out_nd))`
`51`	`51`
`52`	`52`	`def forward(self, x):`
`53`	`53`	`# x: (b, d, n)`