Ok I found out the problem. It is not with weight that has been declared as a Parameter. It is with a map that I define in __init__ function. I need this map to also be put into the same cuda device as weight but I do not want it to be a Parameter.
I tried having these lines
self.w = Parameter(torch.zeros(out_features, in_features))
self.k_map = torch.zeros(self.m+1, device=self.w.device)

Apparently at initialization, the weight is still on the cpu, which causes the map to be on the cpu.
How should I map them both on cuda but not making k_map a Parameter?