Repository: pytorch/ignite

Issue:
`to_onehot` can't be torchscripted
## 🐛 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:
```py
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



Discussion:
Thanks for reporting @ydcjeff ! Let's what can be done in this case...
Here is how it can be implemented.
```python
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) :)
If anyone want to fix it, please go ahead
If not, I can send a PR later.
@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?
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 ?
@Devanshu24 would you like to fix this issue according to the discussion ?
Currently, I am seeing only certain functions from `ignite.utils` which input is tensor.
> @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

Resolving patch:
diff --git a/ignite/utils.py b/ignite/utils.py
--- a/ignite/utils.py
+++ b/ignite/utils.py
@@ -56,8 +56,12 @@ def to_onehot(indices: torch.Tensor, num_classes: int) -> torch.Tensor:
     """Convert a tensor of indices of any shape `(N, ...)` to a
     tensor of one-hot indicators of shape `(N, num_classes, ...) and of type uint8. Output's device is equal to the
     input's device`.
+
+    .. versionchanged:: 0.4.3
+        This functions is now torchscriptable.
     """
-    onehot = torch.zeros(indices.shape[0], num_classes, *indices.shape[1:], dtype=torch.uint8, device=indices.device)
+    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)
 
 


Test patch:
diff --git a/tests/ignite/test_utils.py b/tests/ignite/test_utils.py
--- a/tests/ignite/test_utils.py
+++ b/tests/ignite/test_utils.py
@@ -75,6 +75,29 @@ def test_to_onehot():
     y2 = torch.argmax(y_ohe, dim=1)
     assert y.equal(y2)
 
+    # Test with `TorchScript`
+
+    x = torch.tensor([0, 1, 2, 3])
+
+    # Test the raw `to_onehot` function
+    scripted_to_onehot = torch.jit.script(to_onehot)
+    assert scripted_to_onehot(x, 4).allclose(to_onehot(x, 4))
+
+    # Test inside `torch.nn.Module`
+    class SLP(torch.nn.Module):
+        def __init__(self):
+            super(SLP, self).__init__()
+            self.linear = torch.nn.Linear(4, 1)
+
+        def forward(self, x):
+            x = to_onehot(x, 4)
+            return self.linear(x.to(torch.float))
+
+    eager_model = SLP()
+    scripted_model = torch.jit.script(eager_model)
+
+    assert eager_model(x).allclose(scripted_model(x))
+
 
 def test_dist_setup_logger():