fix: error only when sets are different

ruflab · Nov 14, 2024 · 9c5f208 · 9c5f208
1 parent 486d66a
commit 9c5f208
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/simple_shapes_dataset/data_module.py b/simple_shapes_dataset/data_module.py
@@ -52,7 +52,7 @@ def __init__(
         domains_from_domain_classes = {
             domain_class.base for domain_class in domain_classes
         }
-        if domains_from_props <= domains_from_domain_classes:
+        if domains_from_props < domains_from_domain_classes:
             raise ValueError(
                 f"Domains set in `domain_classes` ({domains_from_domain_classes}) "
                 "are different from domains set in "

diff --git a/simple_shapes_dataset/pre_process.py b/simple_shapes_dataset/pre_process.py
@@ -88,8 +88,9 @@ def nullify_attribute_rotation(
     attr: Sequence[torch.Tensor],
 ) -> list[torch.Tensor]:
     new_attr = attr[1].clone()
-    new_attr[3] = 0.0
-    new_attr[4] = 1.0
+    angle = torch.zeros_like(new_attr[3])
+    new_attr[3] = torch.cos(angle)
+    new_attr[4] = torch.sin(angle)
     return [attr[0], new_attr, attr[2]]