Cannot change improper torsion multiplicity

We get the following error when we restrict the improper multiplicity to n=1,2.

readout_improper = esp.nn.readout.janossy.JanossyPoolingWithSmirnoffImproper(in_features=units, config=janossy_config, out_features={"k": 2})

Error message

  File "/lila/home/takabak/mambaforge/envs/espaloma/lib/python3.9/site-packages/espaloma/mm/energy.py", line 167, in apply_improper_torsion
    % suffix: esp.mm.torsion.periodic_torsion(
  File "/lila/home/takabak/mambaforge/envs/espaloma/lib/python3.9/site-packages/espaloma/mm/torsion.py", line 35, in periodic_torsion
    out = 0.5 * esp.mm.functional.periodic(
  File "/lila/home/takabak/mambaforge/envs/espaloma/lib/python3.9/site-packages/espaloma/mm/functional.py", line 182, in periodic
    torch.nn.functional.relu(k) * (cos_n_theta_minus_phases + 1.0)
RuntimeError: The size of tensor a (2) must match the size of tensor b (6) at non-singleton dimension 2

choderalab / espaloma

Cannot change improper torsion multiplicity #156