`to_onehot` can't be torchscripted

[ydcjeff]

🐛 Bug description

I need to do one hot conversion and I use ignite.utils.to_onehot, but when I try to torchscript the model, it occurs a RuntimeError. Stacktrace:

RuntimeError: 
cannot statically infer the expected size of a list in this context:
  File "/usr/local/lib/python3.6/dist-packages/ignite/utils.py", line 59
    input's device`.
    """
    onehot = torch.zeros(indices.shape[0], num_classes, *indices.shape[1:], dtype=torch.uint8, device=indices.device)
                                                        ~~~~~~~~~~~~~~~~~ <--- HERE
    return onehot.scatter_(1, indices.unsqueeze(1), 1)
'to_onehot' is being compiled since it was called from 'SLP.forward'
  File "<ipython-input-3-0f76d1651906>", line 12
    def forward(self, x):
        x = to_onehot(x, 10)
        ~~~~~~~~~~~~~~~~~~~ <--- HERE
        return self.l1(x)

Environment

Colab

Reproducible Link

https://colab.research.google.com/drive/13mWdz9sGXvAHgKd3vtp5I6Y-gZDRlNAd?usp=sharing

[vfdev-5]

Thanks for reporting @ydcjeff ! Let's what can be done in this case...

[vfdev-5]

Here is how it can be implemented.

def to_onehot(indices: torch.Tensor, num_classes: int) -> torch.Tensor:
    new_shape = (indices.shape[0], num_classes) + indices.shape[1:]
    onehot = torch.zeros(new_shape, dtype=torch.uint8, device=indices.device)
    return onehot.scatter_(1, indices.unsqueeze(1), 1)

@ydcjeff if you'd like to fix it, please send a PR (and if with a test it would be awesome) :)

[ydcjeff]

If anyone want to fix it, please go ahead
If not, I can send a PR later.

[ydcjeff]

@vfdev-5 we should start adding jit tests for the functions or classes that could be used inside a nn.Module.
What do you think?

[vfdev-5]

We can start also doing that. Thanks for the suggestion! I'm wondering if there is a lot of functions like that which could be used inside nn.Module ?

[vfdev-5]

@Devanshu24 would you like to fix this issue according to the discussion ?

[ydcjeff]

Currently, I am seeing only certain functions from ignite.utils which input is tensor.

[Devanshu24]

@Devanshu24 would you like to fix this issue according to the discussion ?

I would love to, I will go through your conversion in a little more detail (have gone through once). I think I will be able to work on it