Fix index of load_pre_saved_latent for ood indicies

bdvllrs · Nov 21, 2023 · b62ea36 · b62ea36
1 parent 1a8a948
commit b62ea36
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/bim_gw/datasets/simple_shapes/datasets.py b/bim_gw/datasets/simple_shapes/datasets.py
@@ -215,7 +215,10 @@ def use_pre_saved_latents(
                             Dict[AvailableDomains, str], pre_saved_latent_path
                         ),
                         domain_key,
-                        self.ids[self.in_dist_labels.shape[0] :],
+                        [
+                            self.ood_idx[k]
+                            for k in self.ids[self.in_dist_labels.shape[0] :]
+                        ],
                     )
                     for k in range(len(data)):
                         self.pre_saved_data[domain_key][k] = np.concatenate(