Expected device cuda:0 and dtype Float but got device cpu and dtype Float

1

Hi all,

I’m working on a project and I’d like to use hook to achieve the goal which is only update some weight tensor in the training.

However, I encountered a problem and tried my best with no result. Below is my codes:

 def hook(self, model, inputs):
        with torch.no_grad():
            print(model.weight.cuda().data.dtype)
            print(self.sparse_mask[self.type[model]].dtype)
            #model.weight.data = model.weight.data * self.sparse_mask[self.type[model]]
            model.weight.cuda().data = model.weight.cuda().data * self.sparse_mask[self.type[model]]
 
    def register_hook(self,module):
        self.handle = module.register_forward_pre_hook(self.hook)

And I got this output

torch.float32
torch.float32

Traceback (most recent call last):
  File "<input>", line 1, in <module>
  File ".pycharm_helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File ".pycharm_helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/sparse_processing.py", line 352, in <module>
    jsegnet(input_data)
  File "/.conda/envs/mengdietao/lib/python3.7/site-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "jsegnet.py", line 199, in forward
    h = self.block1(x)   #[1, 32, 160, 160]
  File "/.conda/envs/mengdietao/lib/python3.7/site-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/jsegnet.py", line 75, in forward
    h = self.conv_a(x)
  File "/.conda/envs/mengdietao/lib/python3.7/site-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/tangshan/tangshan_data/image_segmentation/libs/models/jsegnet.py", line 51, in forward
    return super(_ConvBatchNormReLU, self).forward(x)
  File "/.conda/envs/mengdietao/lib/python3.7/site-packages/torch/nn/modules/container.py", line 92, in forward
    input = module(input)
  File "/.conda/envs/mengdietao/lib/python3.7/site-packages/torch/nn/modules/module.py", line 539, in __call__
    result = hook(self, input)
  File "/sparse_processing.py", line 229, in hook
    model.weight.cuda().data = model.weight.cuda().data * self.sparse_mask[self.type[model]]

RuntimeError: expected device cuda:0 and dtype Float but got device cpu and dtype Float

Both dtype of model.weight.cuda().data and self.sparse_mask[self.type[model]] are torch.float32, my input_data is data = data.cuda().float()

data = data.cuda().float()

I’d like to keep it calculate on GPU but I can’t figure out how to fix that problem.

Any response will be really appreciated!!

2

I GOT THE SOLUTION!

 def hook(self, model, inputs):
        with torch.no_grad():
       
            model.weight.cpu().data = model.weight.cpu().data * self.sparse_mask[self.type[model]]

or

 def hook(self, model, inputs):
        with torch.no_grad():
       
            model.weight.data = model.weight.data * self.sparse_mask[self.type[model]].cuda()