🐛[BUG]: How to use the Stan activation function

[YouqiongLiu]

Version

nvidia-modulus： 0.4.0 nvidia-modulus.sym：1.3.0

On which installation method(s) does this occur?

Pip

Describe the issue

Using the Stan activation function according to the URL: https://github.com/NVIDIA/modulus-sym/blob/main/examples/helmholtz/conf/config_stan , an error message appears. Please provide an example of Stan activation, such as examples/helmholtz/helmholtz_stan.py mentioned in the instructions for use.

Minimum reproducible example

No response

Relevant log output

def forward(self, x):
        if x.shape[-1] != self.beta.shape[-1]:
            raise ValueError(
            ~~~~~~~~~~~~~~~~~
                f"The last dimension of the input must be equal to the dimension of Stan parameters. Got inputs: {x.shape}, params: {self.beta.shape}"
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            )
        return torch.tanh(x) * (1.0 + self.beta * x)
builtins.ValueError: The last dimension of the input must be equal to the dimension of Stan parameters. Got inputs: [2000, 256], params: [4]

Environment details

No response

Other/Misc.

No response