Warning: 'atom_types' is missing in single.xyz. Assuming all atoms are type=0
Traceback (most recent call last):
  File "/run/media/giese/timext14tb1/allegro/allegro_tutorial/../test/run_nequip.py", line 150, in <module>
    ene,frc = eval_allegro_model(model,metadata,crds,atypes)
  File "/run/media/giese/timext14tb1/allegro/allegro_tutorial/../test/run_nequip.py", line 75, in eval_allegro_model
    res = model.forward(inp)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/nequip/nn/_grad_output.py", line 32, in forward
        _7 = torch.append(_3, data[k])
      func0 = self.func
      data0 = (func0).forward(data, )
               ~~~~~~~~~~~~~~ <--- HERE
      of = self.of
      _8 = [torch.sum(data0[of])]
  File "code/__torch__/nequip/nn/_graph_mixin.py", line 28, in forward
    input1 = (radial_basis).forward(input0, )
    input2 = (spharm).forward(input1, )
    input3 = (allegro).forward(input2, )
              ~~~~~~~~~~~~~~~~ <--- HERE
    input4 = (edge_eng).forward(input3, )
    input5 = (edge_eng_sum).forward(input4, )
  File "code/__torch__/allegro/nn/_allegro.py", line 107, in forward
    _15 = torch.cat(latent_inputs_to_cat, -1)
    _16 = annotate(List[Optional[Tensor]], [prev_mask])
    new_latents = (_00).forward(torch.index(_15, _16), )
                   ~~~~~~~~~~~~ <--- HERE
    _17 = annotate(List[Optional[Tensor]], [active_edges0])
    _18 = torch.unsqueeze(torch.index(cutoff_coeffs, _17), -1)
  File "code/__torch__/allegro/nn/_fc.py", line 12, in forward
    x: Tensor) -> Tensor:
    _forward = self._forward
    return (_forward).forward(x, )
            ~~~~~~~~~~~~~~~~~ <--- HERE
class ScalarMLP(Module):
  __parameters__ = []
  File "code/__torch__/torch/fx/graph_module/___torch_mangle_4.py", line 17, in forward
    _weight_0 = self._weight_0
    mul = torch.mul(_weight_0, 0.31622776601683794)
    matmul = torch.matmul(x, mul)
             ~~~~~~~~~~~~ <--- HERE
    silu = __torch__.torch.nn.functional.silu(matmul, False, )
    _weight_1 = self._weight_1

Traceback of TorchScript, original code (most recent call last):
  File "/home/giese/devel/envs/python3.9/local/conda/lib/python3.9/site-packages/nequip/nn/_grad_output.py", line 85, in forward
            wrt_tensors.append(data[k])
        # run func
        data = self.func(data)
               ~~~~~~~~~ <--- HERE
        # Get grads
        grads = torch.autograd.grad(
  File "/home/giese/devel/envs/python3.9/local/conda/lib/python3.9/site-packages/nequip/nn/_graph_mixin.py", line 356, in forward
    def forward(self, input: AtomicDataDict.Type) -> AtomicDataDict.Type:
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/home/giese/devel/envs/python3.9/local/conda/lib/python3.9/site-packages/allegro/nn/_allegro.py", line 500, in forward
    
            # Compute latents
            new_latents = latent(torch.cat(latent_inputs_to_cat, dim=-1)[prev_mask])
                          ~~~~~~ <--- HERE
            # Apply cutoff, which propagates through to everything else
            new_latents = cutoff_coeffs[active_edges].unsqueeze(-1) * new_latents
  File "/home/giese/devel/envs/python3.9/local/conda/lib/python3.9/site-packages/allegro/nn/_fc.py", line 169, in forward
    def forward(self, x):
        return self._forward(x)
               ~~~~~~~~~~~~~ <--- HERE
  File "<eval_with_key>.17", line 7, in forward
    _weight_0 = self._weight_0
    mul = _weight_0 * 0.31622776601683794;  _weight_0 = None
    matmul = torch.matmul(x, mul);  x = mul = None
             ~~~~~~~~~~~~ <--- HERE
    silu = torch.nn.functional.silu(matmul, inplace = False);  matmul = None
    _weight_1 = self._weight_1
RuntimeError: mat1 and mat2 shapes cannot be multiplied (0x8 and 10x32)