The way I reset parameters is the same from the nn.Linear

def reset_parameters(self,W,B=None):
W=torch.nn.init.kaiming_uniform_(W, a=math.sqrt(3))
if B is not None:
fan_in, _ = torch.nn.init.calculate_fan_in_and_fan_out(W)
bound = 1 / math.sqrt(fan_in)
B=torch.nn.init.uniform(B, -bound, bound)
return W,B
return W

I will try without the reset_parameters and as soon as my gpus are free, thank you.