jxx123 opened a new issue #6050:
URL: https://github.com/apache/incubator-tvm/issues/6050
Script to reproduce the issue:
```python
import torch
from tvm import relay
class Foo(torch.nn.Module):
def __init__(self, bar):
super(Foo, self).__init__()
self.bar = bar
def forward(self, x):
return self.bar + x
foo = Foo(torch.cuda.FloatTensor([[1.0, 2.0], [1.0, 2.0]]))
x = torch.cuda.FloatTensor([[1.0, 2.0], [1.0, 2.0]])
traced_foo = torch.jit.trace_module(foo, {'forward': x})
shape_list = [('input0', (2, 2))]
module, params = relay.frontend.from_pytorch(traced_foo, shape_list)
```
Fix is quite straightforward. Just need to convert the tensor to cpu when
the tensor is in cuda at line
https://github.com/jxx123/incubator-tvm/blob/master/python/tvm/relay/frontend/pytorch.py#L2174.
Will submit a PR soon.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
