I am using pretrained Densenet121 from torch vision within my custom nn.module. However, the titled error happens even if I use .cuda(). I tried using .cuda within the init function of CustomNet but it still gives the same error. Please assume I would remove the last layer and replace the classifier of DenseNet with an identity layer before I add the fully-connected layers in the custom net. Thank you very much for your help in advance
class DenseNet121(nn.Module):
    """Model modified.
    The architecture of our model is the same as standard DenseNet121
    except the classifier layer which has an additional sigmoid function.
    """
    def __init__(self, out_size):
        super(DenseNet121, self).__init__()
        self.densenet121 = torchvision.models.densenet121(pretrained=True)
        num_ftrs = self.densenet121.classifier.in_features
        self.in_features = num_ftrs
        self.densenet121.classifier = nn.Sequential(
            nn.Linear(num_ftrs, out_size),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.densenet121(x)
        return x

class CustomNet(nn.Module):
    def __init__(self, dense_net, num_ftrs):
        super(CustomNet, self).__init__()
        self.dense_net = dense_net
        #Fully connected layer1
        self.fc1 = nn.Linear(num_ftrs, 512)
        self.fc2 = nn.Linear(512, 512)
        self.output = nn.Linear(512, 9)
        self.leakyRelu = nn.LeakyReLU(0.01)
        
    def forward(self, x):
        features = self.dense_net(x)
        fc1_output = self.leakyRelu(self.fc1(features ))
        fc2_output = self.leakyRelu(self.fc2(fc1_output))
        #Softmax?
        output = self.output(fc2_output)
        return output


Edit:
Apparently it’s within the conv layers of densenet121, not sure how to solve this