Thank you for your answer @tom
Basically I moved my model and the image and label to the chosen device which is set as an input of my main file. Now I have methods in my model where I create intermediate tensors (torch.zeros or torch.ones for example) and want to allocate them on the chosen device of the model. Unfortunately, these tensors are not parameters and moving the model on the chosen device doesn‘t create these tensors automatically on the device which the model is allocated to. So I have to use the device as an input of the model so that the methods can use it as an input. Is there a better way to do this other than having the device object as an input in methods where tensors are created?