I am writing a class for a restricted boltzmann machine using PyTorch. I am using the weights and the biases within a call to LBFGS optimizer, so I’d like to have all the weights and biases as Parameters. However, when I try to initialize them as cuda tensors, I am getting the following error:
TypeError: cannot assign 'torch.cuda.FloatTensor' as parameter 'W' (torch.nn.Parameter or None expected) 

Here is my code:
class RBM(nn.Module):
    
    def __init__(self,
                    visible_units   = 256,
                    hidden_units    = 64,
                    epsilonw        = 0.1, #learning rate for weights
                    epsilonvb       = 0.1, #learning rate for visible unit biases
                    epsilonhb       = 0.1, #learning rate for hidden unit biases
                    weightcost      = 0.0002,
                    initialmomentum = 0.5,
                    finalmomentum   = 0.9,
                    use_gpu = False
                    ):
        super(RBM,self).__init__()
        
        self.desc               = "RBM"
        self.visible_units      = visible_units
        self.hidden_units       = hidden_units
        self.epsilonw           = epsilonw
        self.epsilonvb          = epsilonvb
        self.epsilonhb          = epsilonhb
        self.weightcost         = weightcost
        self.initialmomentum    = initialmomentum
        self.finalmomentum      = finalmomentum
        self.use_gpu            = use_gpu
        self.device             = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        self.h_bias = nn.Parameter(torch.zeros(self.hidden_units)) #hidden layer bias
        self.v_bias = nn.Parameter(torch.zeros(self.visible_units)) #visible layer bias
           
        if self.use_gpu:
            self.W      = nn.Parameter(0.1*torch.randn(self.visible_units, self.hidden_units), 
                                       device=self.device)
            self.v_bias = nn.Parameter(torch.zeros(self.hidden_units),
                                       device=self.device)
            self.h_bias = nn.Parameter(torch.zeros(self.visible_units),
                                       device=self.device)
        else:
            self.W      = nn.Parameter(0.1*torch.randn(self.visible_units, self.hidden_units))
            self.h_bias = nn.Parameter(torch.zeros(self.hidden_units)) #hidden layer bias
            self.v_bias = nn.Parameter(torch.zeros(self.visible_units)) #visible layer bias

The full traceback for the call is:
from RBM import RBM
rbm = RBM(use_gpu=1)
Traceback (most recent call last):  
File "<stdin>", line 1, in <module>
File "/home/deep_autoencoder/RBM.py", line 54, in __init__                                                   
    self.W = nn.Parameter(0.1*torch.randn(self.visible_units, self.hidden_units),
File "/home/anaconda3/envs/torch_env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 792, in __setattr__                          
    raise TypeError("cannot assign '{}' as parameter '{}' "  
TypeError: cannot assign 'torch.cuda.FloatTensor' as parameter 'W' (torch.nn.Parameter or None expected) 

I’m not sure what I’m doing wrong here. The code I’ve written seems to reflect the best advice I’ve seen in this forum.