Merge pull request google-research#262 from RerRayne/patch-1

Fix wrong indent in point tracking + correct batching order.
cdoersch · Jul 26, 2022 · b42f616 · b42f616
2 parents 509bfa2 + d864237
commit b42f616
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/challenges/point_tracking/dataset.py b/challenges/point_tracking/dataset.py
@@ -681,17 +681,17 @@ def create_point_tracking_dataset(
   ds = ds[split]
   if repeat:
     ds = ds.repeat()
-    ds = ds.map(
-        functools.partial(
-            add_tracks,
-            train_size=train_size,
-            vflip=vflip,
-            random_crop=random_crop),
-        num_parallel_calls=2)
+  ds = ds.map(
+      functools.partial(
+          add_tracks,
+          train_size=train_size,
+          vflip=vflip,
+          random_crop=random_crop),
+      num_parallel_calls=2)
   if shuffle_buffer_size is not None:
     ds = ds.shuffle(shuffle_buffer_size)
 
-  for bs in batch_dims:
+  for bs in batch_dims[::-1]:
     ds = ds.batch(bs)
 
   return ds