Input tensor is not an XLA tensor: CPUFloatType while using crf.decode function about aws-neuron-sdk

[tombettany]

We had a similar problem when trying to run torch_neuronx.trace() and it was because the output of our model had the device set for cpu rather than an XLA device. Solution we found was just to set FloatTensor(a, device=torch_xla.core.xla_model.xla_device()) when creating the tensor to set the device type. Alternatively you can run .to(torch_xla.core.xla_model.xla_device()) on the output:

import torch_xla.core.xla_model as xm

def tags(output , mask):
    return model.crf.decode(output, mask).to(xm.xla_device())

...

Presumably you could set the device type on the input tensor instead but we did have issues with this on our model.

Hopefully this gets you a little bit further towards solving the problem.

[PrateekAg1511]

@tombettany Thanks!

I tried this but then got the following warning:

/usr/local/lib/python3.10/site-packages/torch_neuronx/xla_impl/hlo_conversion.py:143: UserWarning: Received an input tensor that was unused. Tensor will be ignored. (index=0, shape=torch.Size([1, 60, 184]), dtype=torch.float32)
  warnings.warn(
/usr/local/lib/python3.10/site-packages/torch_neuronx/xla_impl/hlo_conversion.py:143: UserWarning: Received an input tensor that was unused. Tensor will be ignored. (index=1, shape=torch.Size([1, 60]), dtype=torch.uint8)
  warnings.warn(

Now the traced model is giving the same output for every input that it gets.

[jluntamazon]

Would you be able to give more details about the model you are trying to trace? If there is a minimal open source reproduction of the error you are encountering, we can try to help you solve the problem.

The warning you are running into indicates that the output of the model does not appear to depend on the inputs. This can happen when the output is calculated entirely based on tensors which are newly constructed within the forward function. This likely happens due to the implementation of the model.crf.decode method.

[PrateekAg1511]

@jluntamazon Here is the minimal open source reproduction of the error:

import torch
from torchcrf import CRF
num_tags = 184
model = CRF(num_tags)

emissions = torch.rand([1,60,184])
mask = torch.ones([1,60], dtype=torch.uint8)

def decode_fn(emissions , mask):
a = model.decode(emissions , mask)
a = torch.Tensor(a)
a = a.to(xm.xla_device())
return (a)

inputs_crf = emissions , mask

trace_crf = torch_neuronx.trace(decode_fn , inputs_crf)

After running trace, I get the waring message for both inputs.

"/aws_neuron_venv_pytorch/lib64/python3.9/site-packages/torch_neuronx/xla_impl/hlo_conversion.py:144: UserWarning: Received an input tensor that was unused. Tensor will be ignored. (index=0, shape=torch.Size([1, 60, 184]), dtype=torch.float32)
warnings.warn(
//aws_neuron_venv_pytorch/lib64/python3.9/site-packages/torch_neuronx/xla_impl/hlo_conversion.py:144: UserWarning: Received an input tensor that was unused. Tensor will be ignored. (index=1, shape=torch.Size([1, 60]), dtype=torch.uint8)
warnings.warn( "