Right. So as described in my original post, softmax function has also been moved to CUDA. I briefly reviewed the source code for nn.Softmax() and saw these lines:
class Softmax(Module):
  
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -> None:
        super(Softmax, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -> Tensor:
        return F.softmax(input, self.dim, _stacklevel=5)

    def extra_repr(self) -> str:
        return 'dim={dim}'.format(dim=self.dim)

I am not sure why def__setstate__ is used and what state is it setting? How does it receive the state? Also, I noticed the module/function that inherits nn.Module class usually gets to move to CUDA. Any thoughts?