Select head of multihead mace model

[Felixrccs]

Is it possible to select a specific head of a mutlihead Mace model?

[thomasloux]

Do you have an example?

[Felixrccs]

Sadly I dont have a model that I can hand out. But I solved the issue: One can pass a tensor with the systeme size and the head ID through to the model. In this case head 1.

 # Head for whole batch
        head = torch.ones(self.n_systems, device=self.device, dtype=torch.int32)

        # Get model output
        out = self.model(
            dict(
                ptr=self.ptr,
                node_attrs=self.node_attrs,
                batch=sim_state.system_idx,
                pbc=sim_state.pbc,
                cell=sim_state.row_vector_cell,
                positions=sim_state.positions,
                edge_index=edge_index,
                unit_shifts=unit_shifts,
                shifts=shifts,
                head=head
            ),
            compute_force=self.compute_forces,
            compute_stress=self.compute_stress,
        )