Error of using ssp as activation function

[Skyinner]

I use ssp as activation function. However, the following error exists. I found that removing @torch.jit.script of ssp function in nequip can help run smoothly.

Traceback (most recent call last):
  File "/Users/shipengjie/anaconda3/lib/python3.10/site-packages/nequip/utils/auto_init.py", line 232, in instantiate
    instance = builder(positional_args, final_optional_args)
  File "/Users/shipengjie/anaconda3/lib/python3.10/site-packages/allegro/nn/allegro.py", line 316, in __init__
    two_body_latent(
  File "/Users/shipengjie/anaconda3/lib/python3.10/site-packages/allegro/nn/_fc.py", line 153, in __init__
    features = nonlinearity(features)
RuntimeError: ShiftedSoftPlus() Expected a value of type 'Tensor (inferred)' for argument 'x' but instead found type 'Proxy'.
Inferred 'x' to be of type 'Tensor' because it was not annotated with an explicit type.
Position: 0
Value: Proxy(matmul)
Declaration: ShiftedSoftPlus(Tensor x) -> (Tensor)
Cast error details: Unable to cast Proxy(matmul) to Tensor

The system I am running is macos.

[Linux-cpp-lisp]

Hi @Skyinner ,

Out of curiosity, why do you need to use SSP instead of SiLU? Have you found better accuracy on some system?

This is I think because SSP needs to be rewritten as a function rather than a module...

[Skyinner]

Hi @Linux-cpp-lisp , I just tested the SSP activation function and found the error. After using SSP, I did not find any significant performance advantages.

[Linux-cpp-lisp]

Got it, I'll keep this issue open since it is techincaly a bug, but hopefully you can just move forward with your work with SiLU. Thanks!