How about the following?
def cuda(self, device_id=None):
        """Moves all model parameters and buffers to the GPU.

        Arguments:
            device_id (int, optional): if specified, all parameters will be
                copied to that device
        """
        self._cuda = Ture
        self._device_id = device_id
        return self._apply(lambda t: t.cuda(device_id))

def cpu(self, device_id=None):
        """Moves all model parameters and buffers to the CPU."""
        self._cuda = False
        return self._apply(lambda t: t.cpu())

def get_parameter(self, name, shape):
       if not name in self._parameters.keys():
            self._parameters[name] = nn.Parameter(torch.randn(shape))
       var = self._parameters[name]      
       fn = lambda t: t.cuda(self._device_id) if self._cuda else lambda t: t.cpu()
       var.data = fn(var.data)
       if var.grad is not None: 
           var.grad.data = fn(var.grad.data)
       return var