dataset fixing

2020-06-19 15:37:43 +02:00
parent 3f8122484b
commit aea34de964
2 changed files with 6 additions and 4 deletions
--- a/point_toolset/point_io.py
+++ b/point_toolset/point_io.py
@ -11,9 +11,9 @@ class BatchToData(object):
        # Convert to torch_geometric.data.Data type

        batch_pos = batch_dict['pos']
-        batch_norm = batch_dict['norm']
-        batch_y = batch_dict['y']
-        batch_y_c = batch_dict['y_c']
+        batch_norm = batch_dict.get('norm', None)
+        batch_y = batch_dict.get('y', None)
+        batch_y_c = batch_dict.get('y_c', None)

        batch_size, num_points, _ = batch_pos.shape  # (batch_size, num_points, 3)