Hi,
I tried adding torch.Variable to my custom (nn.module) and add it to optimizer group.
Variable name is ‘self.scale’
Code is like below.
class K32_BatchNet(nn.Module):
     def __init__(self):
         super(K32_BatchNet, self).__init__()
         self.layer1 = nn.Sequential(
             nn.Conv2d(1, 32, kernel_size=5, padding=2),
             nn.BatchNorm2d(32))
         self.scale = Variable(torch.ones(1).cuda(), requires_grad=True)
         self.layer2 = nn.Sequential(
             nn.Conv2d(32, 64, kernel_size=5, padding=2),
             nn.BatchNorm2d(64))
         self.layer3 = nn.Sequential(
             nn.Conv2d(64, 128, kernel_size=5, padding=2),
             nn.BatchNorm2d(128))
         self.conv3_drop = nn.Dropout2d()
         self.avgPool2d = nn.AvgPool2d(3, stride=2)
         self.fc1 = nn.Linear(1152,160)
         self.fc2 = nn.Linear(160, 10)
 
     def forward(self, x):
        out = F.relu(self.layer1(x))
         out = CustomFunction()(out, self.scale)
        out = F.max_pool2d(out, 2)

Then, in main function,
    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=5e-4)
    optimizer.add_param_group({"params": model.scale})

In train function which is called in the main function,
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

However, it shows below error.
Traceback (most recent call last):
  File "k32_train.py", line 298, in <module>
    main()
  File "k32_train.py", line 93, in main
    train(train_loader, model, optimizer, epoch)
  File "k32_train.py", line 231, in train
    optimizer.step()
  File "/SSD/root/.local/lib/python2.7/site-packages/torch/optim/sgd.py", line 93, in step
    d_p.add_(weight_decay, p.data)
RuntimeError: expand(torch.cuda.FloatTensor{[1]}, size=[]): the number of sizes provided (0) must be greater or equal to the number of dimensions in the tensor (1)

Please help me to get out of this problem.
Thanks!