Can I use dictionary in an extension of torch.autograd.Function when using GPU?
I have the following code in an extension of torch.autograd.Function: where y has type LongTensor, self.N and self.C have type long. it works fine if I use CPU, but got the following error if GPU is used: is this error because of I am using dictionary and GPU simultaneously or I did something else wrong? Thank you very much!
lili.ece.gwu
apaszke
I really appreciate your answer, but the next problem is how I can know the function is run in GPU or CPU, is there any flag indicates it? I tried to pass a bool argument from outside of the function but another error raises. Is it because of the torch.autograd.Function accepts Variable only? if so, how can I determine the function is run in GPU or not? Thank you!
lili.ece.gwu
apaszke
Many thanks!! it works now!!
lili.ece.gwu