Force a tensor to be on cpu

I have created a model which inherits nn.Module and I have transferred to model.cuda() . While passing the parameters in forward function i want one argument to be in cpu but since I have used model.cuda() all arguments are in turn being converted to gpu i think. How do i force a tensor to live in cpu ?

Have you tried Tensor.cpu() for the parameter which you want to run on cpu?

Tensor.cpu() will transfer to cpu but the point of forcing the tensor in cpu is because my tensor is a big matrix and transferring to gpu and then to cpu is not necessary.

You can partially choose cpu or gpu for each weight. Did you mean this?

model = models.resnet18()
model.cuda()
model.layer1[0].conv1.weight.type()
'torch.cuda.FloatTensor'
model.layer1[0].conv2.weight.type()
'torch.cuda.FloatTensor'
model.layer1[0].conv1.cpu()
Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
model.layer1[0].conv1.weight.type()
'torch.FloatTensor'
model.layer1[0].conv2.weight.type()
'torch.cuda.FloatTensor'
class Model(nn.Module):
    def __init__(self):
         pass
    def forward(self, inputs1, inputs2, inputs3):
         # inputs1, inputs2, inputs3 are already in gpu.
         pass
model = Model()
model = torch.nn.DataParallel(model)
model = model.cuda()

Here the inputs1, inputs2, inputs3 are in gpu even though i pass them as just tensor with cpu because of the model.cuda(). The inputs2 is so big that crashes the program with cuda out of memory . Is there a way to force inputs2 to be in cpu without converting them to cpu in the forward function. ?

1 Like
    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            return self.module(*inputs, **kwargs)
        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            return self.module(*inputs[0], **kwargs[0])
        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = self.parallel_apply(replicas, inputs, kwargs)
        return self.gather(outputs, self.output_device)

The self.scatter in the forward function of dataparallel is trying to port all arguments into the gpu but i dont want inputs2 to be in gpu

Even if you do something to force inputs2 to cpu then you cannot do operations with that inputs2 because it becomes torch.FloatTensor and your model parameters will be torch.cuda.FloatTensor which are incompatible