Fix symmetry issues (#126)

Addressing #115

Author: Ubiquinone-dot

models/rfd3/src/rfd3/engine.py

@@ -5,7 +5,7 @@
 from dataclasses import dataclass, field
 from os import PathLike
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 
 import torch
 import yaml
@@ -378,7 +378,7 @@ def _canonicalize_inputs(
 
     def _multiply_specifications(
         self, inputs: Dict[str, dict | DesignInputSpecification], n_batches=None
-    ) -> Dict[str, Dict[str, Any]]:
+    ) -> Dict[str, dict | DesignInputSpecification]:
         # Find existing example IDS in output directory
         if exists(self.out_dir):
             existing_example_ids = set(

models/rfd3/src/rfd3/utils/inference.py

@@ -373,12 +373,13 @@ def ensure_inference_sampler_matches_design_spec(
         design_spec: Design specification dictionary
         inference_sampler: Inference sampler dictionary
     """
-    has_symmetry_specification = [
-        True
-        if "symmetry" in item.keys() and item.get("symmetry") is not None
-        else False
-        for item in design_spec.values()
-    ]
+    has_symmetry_specification = []
+    for item in design_spec.values():
+        if hasattr(item, "symmetry"):
+            has_symmetry = item.symmetry is not None
+        else:
+            has_symmetry = "symmetry" in item and item.get("symmetry") is not None
+        has_symmetry_specification.append(has_symmetry)
     if any(has_symmetry_specification):
         if (
             inference_sampler is None