diff --git a/src/metatrain/experimental/flashmd/model.py b/src/metatrain/experimental/flashmd/model.py index d20826f517..08711d1d55 100644 --- a/src/metatrain/experimental/flashmd/model.py +++ b/src/metatrain/experimental/flashmd/model.py @@ -588,7 +588,7 @@ def forward( ) ) else: - output_blocks.append(b) + output_blocks.append(b.copy()) return_dict[name] = TensorMap( return_dict[name].keys, output_blocks ) diff --git a/src/metatrain/experimental/flashmd_symplectic/model.py b/src/metatrain/experimental/flashmd_symplectic/model.py index 959e288b82..d56005073e 100644 --- a/src/metatrain/experimental/flashmd_symplectic/model.py +++ b/src/metatrain/experimental/flashmd_symplectic/model.py @@ -643,7 +643,7 @@ def forward( ) ) else: - output_blocks.append(b) + output_blocks.append(b.copy()) return_dict[name] = TensorMap( return_dict[name].keys, output_blocks ) diff --git a/src/metatrain/experimental/mace/model.py b/src/metatrain/experimental/mace/model.py index 31ce13d763..758af4c0cf 100644 --- a/src/metatrain/experimental/mace/model.py +++ b/src/metatrain/experimental/mace/model.py @@ -528,7 +528,7 @@ def add_additive_contributions( ) ) else: - output_blocks.append(b) + output_blocks.append(b.copy()) values[name] = TensorMap(values[name].keys, output_blocks) def supported_outputs(self) -> Dict[str, ModelOutput]: diff --git a/src/metatrain/experimental/phace/model.py b/src/metatrain/experimental/phace/model.py index 6982fb882d..b44cade121 100644 --- a/src/metatrain/experimental/phace/model.py +++ b/src/metatrain/experimental/phace/model.py @@ -513,7 +513,7 @@ def forward( ) ) else: - output_blocks.append(b) + output_blocks.append(b.copy()) return_dict[name] = TensorMap(return_dict[name].keys, output_blocks) # For atomic basis targets, sparsify to create blocks with "atom_type" diff --git a/src/metatrain/gap/model.py b/src/metatrain/gap/model.py index 51b5815d80..ee83451b9d 100644 --- a/src/metatrain/gap/model.py +++ b/src/metatrain/gap/model.py @@ -243,9 +243,17 @@ def forward( for idx_key in range(len(self._species_labels)): key = self._species_labels.entry(idx_key) if soap_features.keys.position(key) is not None: - new_blocks.append(soap_features.block(key)) + block = soap_features.block(key) + new_blocks.append( + TensorBlock( + values=block.values, + samples=block.samples, + components=block.components, + properties=block.properties, + ) + ) else: - new_blocks.append(dummyblock) + new_blocks.append(dummyblock.copy()) soap_features = TensorMap(keys=self._species_labels, blocks=new_blocks) soap_features = soap_features.keys_to_samples("center_type") # here, we move to properties to use metatensor operations to aggregate @@ -254,7 +262,18 @@ def forward( soap_features = soap_features.keys_to_properties( ["neighbor_1_type", "neighbor_2_type"] ) - soap_features = TensorMap(self._keys, soap_features.blocks()) + soap_features = TensorMap( + self._keys, + [ + TensorBlock( + values=block.values, + samples=block.samples, + components=block.components, + properties=block.properties, + ) + for block in soap_features.blocks() + ], + ) output_key = list(outputs.keys())[0] energies = self._subset_of_regressors_torch(soap_features) return_dict = {output_key: energies} diff --git a/src/metatrain/soap_bpnn/model.py b/src/metatrain/soap_bpnn/model.py index 48a175b392..694661122e 100644 --- a/src/metatrain/soap_bpnn/model.py +++ b/src/metatrain/soap_bpnn/model.py @@ -865,7 +865,7 @@ def forward( ) ) else: - output_blocks.append(b) + output_blocks.append(b.copy()) return_dict[name] = TensorMap(return_dict[name].keys, output_blocks) # For atomic basis targets, sparsify to create blocks with "atom_type" diff --git a/src/metatrain/utils/data/target_info.py b/src/metatrain/utils/data/target_info.py index 9ba103ff52..47feaea5ea 100644 --- a/src/metatrain/utils/data/target_info.py +++ b/src/metatrain/utils/data/target_info.py @@ -758,7 +758,10 @@ def is_auxiliary_output(name: str) -> bool: :return: `True` if the target is an auxiliary output, `False` otherwise. """ is_auxiliary = ( - name == "features" or name == "energy_ensemble" or name.startswith("mtt::aux::") + name == "features" + or name == "feature" + or name == "energy_ensemble" + or name.startswith("mtt::aux::") ) return is_auxiliary diff --git a/src/metatrain/utils/evaluate_model.py b/src/metatrain/utils/evaluate_model.py index 6a0bb0006b..0308bcb972 100644 --- a/src/metatrain/utils/evaluate_model.py +++ b/src/metatrain/utils/evaluate_model.py @@ -62,7 +62,7 @@ def evaluate_model( energy_targets_that_require_strain_gradients = [] for target_name in targets.keys(): # Check if the target is an energy: - if model_outputs[target_name].quantity == "energy": + if targets[target_name].quantity == "energy": energy_targets.append(target_name) if isinstance(targets[target_name], TargetInfo): # Check if the energy requires gradients: diff --git a/src/metatrain/utils/scaler/_base_scaler.py b/src/metatrain/utils/scaler/_base_scaler.py index 0584fd17d0..ad69e133b7 100644 --- a/src/metatrain/utils/scaler/_base_scaler.py +++ b/src/metatrain/utils/scaler/_base_scaler.py @@ -831,11 +831,11 @@ def _set_fixed_weights( self.scales[target_name] = TensorMap( self.Y2[target_name].keys.to(device=block.values.device), - [block], + [block.copy()], ) self.per_target_scales[target_name] = TensorMap( self.Y2[target_name].keys.to(device=block.values.device), - [block], + [block.copy()], ) def _sync_device_dtype(self, device: torch.device, dtype: torch.dtype) -> None: diff --git a/src/metatrain/utils/testing/output.py b/src/metatrain/utils/testing/output.py index 967e3e6fee..68a911cf4f 100644 --- a/src/metatrain/utils/testing/output.py +++ b/src/metatrain/utils/testing/output.py @@ -627,7 +627,7 @@ def test_output_features( features_output_options = ModelOutput( quantity="", - unit="unitless", + unit="", per_atom=per_atom, ) model = model.to(system.positions.dtype) @@ -705,7 +705,7 @@ def test_output_last_layer_features( # last-layer features per atom: ll_output_options = ModelOutput( quantity="", - unit="unitless", + unit="", per_atom=per_atom, ) model = model.to(system.positions.dtype)