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:
sdesrozis: @Devanshu24 Thank you for this PR !

Could you add a test using torchscript ? Thanks !
Devanshu24: > Could you add a test using torchscript ? Thanks !

Sure, I'll do that! :)


vfdev-5: @Devanshu24 could you please advance on this PR in higher priority than the one with deprecated decorator as we'd like to put it into the next release. Thanks !
Devanshu24: Sure, I was learning how to use torchscript. Will try to push a commit today itself :)
Devanshu24: While writing the tests for `TorchScript` I noticed something
The following code gives `True` for all elements of the tensor

```python
from ignite.utils import to_onehot
from torch import nn
import torch


class SLP(nn.Module):
    def __init__(self):
        super(SLP, self).__init__()
        self.linear = nn.Linear(4, 1)

    def forward(self, x):
        x = to_onehot(x, 4)
        return self.linear(x.to(torch.float))

a = torch.tensor([0, 1, 2, 3])
eager_model = SLP()
print(eager_model(a))
script_model = torch.jit.trace(eager_model, a)
print(script_model(a))

torch.eq(eager_model(a), script_model(a))
```
______

However, if I pass a different object (of the same class) to `torch.jit.trace` the outputs of the `eager_model` and `script_model` are totally different (not even close)
See the code below, the output is `False` for all elements of the tensor
```python
from ignite.utils import to_onehot
from torch import nn
import torch


class SLP(nn.Module):
    def __init__(self):
        super(SLP, self).__init__()
        self.linear = nn.Linear(4, 1)

    def forward(self, x):
        x = to_onehot(x, 4)
        return self.linear(x.to(torch.float))

a = torch.tensor([0, 1, 2, 3])
eager_model = SLP()
print(eager_model(a))
eager_model_new = SLP()
script_model = torch.jit.trace(eager_model_new, a)
print(script_model(a))

torch.eq(eager_model(a), script_model(a))
```

I have only started to learn about `TorchScript` so I do not know if this is expected behaviour. It'd be great if someone could guide me on this _or_ direct me to the correct resource.
vfdev-5: @Devanshu24 maybe it should depend on how fully-connected weights are randomly initialized. Put `torch.manual_seed(12)` before each model creation and it should be good. By the way, let's use `torch.jit.script` for scripting ops.
So, a basic test without inserting `to_onehot` into nn.Module can be
```python
scripted_to_onehot = torch.jit.script(to_onehot)
x = ...
assert scripted_to_onehot(x).allclose(to_onehot(x))
```
Devanshu24: > @Devanshu24 maybe it should depend on how fully-connected weights are randomly initialized. Put `torch.manual_seed(12)` before each model creation and it should be good.

Ohh right! Thank you!!

> So, a basic test without inserting `to_onehot` into nn.Module can be
> 
> ```python
> scripted_to_onehot = torch.jit.script(to_onehot)
> x = ...
> assert scripted_to_onehot(x).allclose(to_onehot(x))
> ```
Okay, I will use this, thanks! :)

Devanshu24: I have added `TorchScript` based tests for the `to_onehot` function. Please have a look at your convenience :)
vfdev-5: @Devanshu24 I updated the test myself to see if everything is ok and we can merge the PR 
Devanshu24: Thanks @vfdev-5 !
Sorry that  `.trace` call skipped my sight, thanks for catching it! :)
ydcjeff: Looks good, just docs thing
```py
.. versionchanged:: 0.4.3
    This functions is now torchscriptable.
```
cc: @vfdev-5 for version number
vfdev-5: Thanks for the reminder @ydcjeff ! This new practice is not completely intergrated :) Maybe, we can update PR template to recall (if anyone actually read it :)

@Devanshu24 could you please add versionchanged docs mark as suggested. Thanks  
Devanshu24: On a side note, why is the docs preview action not triggered?
vfdev-5: > On a side note, why is the docs preview action not triggered?

I disabled to consume less in building minutes... Will enable it since tomorrow as it will be new month with more available credits :)

Patch:
diff --git a/ignite/utils.py b/ignite/utils.py
index cd00e8e1..91a9b8f4 100644
--- 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
index 781bd267..1b523e86 100644
--- 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():