Merge branch 'main' into move_to_autocast_general_and_amp_test

src/fairchem/core/datasets/base_dataset.py

@@ -109,7 +109,9 @@ def _metadata(self) -> DatasetMetadata:
                 for field in DatasetMetadata._fields
             }
         )
-
+        assert np.issubdtype(
+            metadata.natoms.dtype, np.integer
+        ), f"Metadata natoms must be an integer type! not {metadata.natoms.dtype}"
         assert metadata.natoms.shape[0] == len(
             self
         ), "Loaded metadata and dataset size mismatch."

src/fairchem/core/trainers/ocp_trainer.py

@@ -262,10 +262,8 @@ def _forward(self, batch):
                     ), f"we need to know which property to match the target to, please specify the property field in the task config, current config: {self.output_targets[target_key]}"
                     prop = self.output_targets[target_key]["property"]
                     pred = out[target_key][prop]
-
-            ## TODO: deprecate the following logic?
-            ## Otherwise, assume target property is a derived output of the model. Construct the parent property
-            else:
+            # TODO clean up this logic to reconstruct a tensor from its predicted decomposition
+            elif "decomposition" in self.output_targets[target_key]:
                 _max_rank = 0
                 for subtarget_key in self.output_targets[target_key]["decomposition"]:
                     _max_rank = max(
@@ -293,6 +291,10 @@ def _forward(self, batch):
                     cg_change_mat(_max_rank, self.device),
                     pred_irreps,
                 )
+            else:
+                raise AttributeError(
+                    f"Output target: '{target_key}', not found in model outputs: {list(out.keys())}"
+                )
 
             ### not all models are consistent with the output shape
             ### reshape accordingly: num_atoms_in_batch, -1 or num_systems_in_batch, -1