Issue using spkdeploy

[jap93]

Hi I have trained a model using a recent version of schnetpack (10.5.24). I am using spkdeploy to get a model to use in an alternate program but keep getting the following error

Traceback (most recent call last):
  File "/mnt/Progs/python/schnetpack/src/scripts/spkdeploy", line 49, in <module>
    save_jit_model(model, args.deployed_model_path)
  File "/mnt/Progs/python/schnetpack/src/scripts/spkdeploy", line 31, in save_jit_model
    jit_model = get_jit_model(model)
  File "/mnt/Progs/python/schnetpack/src/scripts/spkdeploy", line 27, in get_jit_model
    return torch.jit.script(model)
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_script.py", line 1338, in script
    return torch.jit._recursive.create_script_module(
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 558, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 631, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_script.py", line 647, in _construct
    init_fn(script_module)
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 607, in init_fn
    scripted = create_script_module_impl(
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 635, in create_script_module_impl
    create_methods_and_properties_from_stubs(
  File "/home/jap93/.local/lib/python3.10/site-packages/torch/jit/_recursive.py", line 467, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(
RuntimeError: 
Module 'SchNet' has no attribute 'charge_embedding' :
  File "/home/jap93/.local/lib/python3.10/site-packages/schnetpack/representation/schnet.py", line 185
            num_batch = len(inputs[structure.idx])
    
            charge_embedding = self.charge_embedding(
                               ~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                x, total_charge, num_batch, idx_m
            )

I am confused why I am getting this error as I have done this on a previous occasion without problem.

John

[jnsLs]

Dear John,

we changed some things in the SchNet class a few weeks ago.
Most likely your error occurs because of that change.
I will provide a fix asap.

Best, Jonas

[jnsLs]

For your particular problem, it seems that you are not using spin_embedding or charge_embedding. Correct?
So if you don't have time to wait for our fix, you can work with the following quick and dirty work around:

Just comment out the following block in the SchNet forward method:

        # add spin and charge embeddings
        if hasattr(self, "activate_charge_spin_embedding") and self.activate_charge_spin_embedding:
            # get tensors from input dictionary
            total_charge = inputs[structure.total_charge]
            spin = inputs[structure.spin_multiplicity]
            idx_m = inputs[structure.idx_m]
            num_batch = len(inputs[structure.idx])

            charge_embedding = self.charge_embedding(
                x, total_charge, num_batch, idx_m
            )
            spin_embedding = self.spin_embedding(
                x, spin, num_batch, idx_m
            )

            # additive combining of nuclear, charge and spin embedding
            x = x + charge_embedding + spin_embedding

In the mean time we will work on a fix.