I’ve defined a submodule which contains parameters I want to penalize. To be specific, I want to penalize the 1st dim of this parameter only!
Here is a demo of my implementation. it works on cpu, while throws RuntimeError on ‘cuda’.
import torch
import numpy as np

class mySubModule(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.my_param = torch.nn.Parameter(
            torch.empty(n, 2).uniform_(0.0, 1.0), requires_grad=True)
        self.register_buffer('mask_choice', torch.tensor([[1.], [0.]]))

    def forward(self, x):
        out = torch.matmul(
            torch.matmul(x, self.my_param),
            self.mask_choice)
        return out

class myModule(torch.nn.Module):
    def __init__(self, n):
        super().__init__()
        self.subModule = mySubModule(n)
        self.linear = torch.nn.Linear(n, 1)

        self.param_need_l1_penalty_case_1 = [self.subModule.my_param]
        #### Here's where Error happens
        self.param_need_l1_penalty_case_2 = [self.subModule.my_param[:,0]]
        ####

    def forward(self, x):
        return self.linear(x) + self.subModule(x)

a tipical useage during training looks like
# demo of trainer
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
np.random.seed(999)

model = myModule(n=2)
model.to(device)

my_data = torch.from_numpy(np.random.random((4,2))).to(torch.float).to(device)
label = torch.from_numpy(np.ones((4,))).to(torch.float).to(device)

pred = model(my_data).squeeze(1)
criteria = torch.nn.MSELoss()
loss = criteria(label, pred)
# add penalty
for param in model.param_need_l1_penalty_case_2:
    loss += 0.1 * torch.norm(param, 1)
print(loss)
loss.backward()

When cuda is available, it raises error saying

RuntimeError: Function AddBackward0 returned an invalid gradient at index 1 - expected type torch.FloatTensor but got torch.cuda.FloatTensor

Any suggestion on how to add l1/l2 penalization in this scenario?
Appreciated~