Can now be trained with normals

2019-08-09 15:37:26 +02:00
parent 92117328ad
commit 167ac4991e
2 changed files with 13 additions and 5 deletions
@@ -148,6 +148,8 @@ class CustomShapeNet(InMemoryDataset):
                    ####################################
                    # This is where you define the keys
                    attr_dict = dict(y=y, pos=points[:, :3 if not self.with_normals else 6])
+                    if not self.with_normals:
+                        attr_dict.update(normals=points[:, 3:6])
                    ####################################
                    if self.collate_per_element:
                        data = Data(**attr_dict)