restructured

2020-03-11 19:34:05 +01:00
parent 1f4edae95c
commit 7b795c2f7b
3 changed files with 49 additions and 69 deletions
--- a/datasets/trajectory_dataset.py
+++ b/datasets/trajectory_dataset.py
@@ -84,7 +84,7 @@ class TrajData(object):
    def map_shapes_max(self):
        shapes = self.map_shapes
        shape_list = list(map(max, zip(*shapes)))
-        if self.mode == 'all_in_map':
+        if self.mode in ['separated_arrays', 'all_in_map']:
            shape_list[0] += 2
        return shape_list