Don’t pass the device argument to SkipeRelu, but use the .device attribute from e.g. the incoming activation tensor, if needed.
E.g. you could use:
     def reset_vmem(self, device):
         self.vmem = torch.zeros(*self.sz, device=device, requires_grad=True)

and call it via self.reset_vmem(x.device) where x would be the input tensor.