Fix YoloNAS on cuda

[Louis-Dupont]

Reproduce bug

import torch
from super_gradients.common.object_names import Models
from super_gradients.training import models

# Note that currently only YoloX, PPYoloE and YOLO-NAS are supported.
model = models.get(Models.YOLO_NAS_L, pretrained_weights="coco")

# We want to use cuda if available to speed up inference.
model = model.to("cuda" if torch.cuda.is_available() else "cpu")

IMAGES = [
    "../../../../documentation/source/images/examples/countryside.jpg",
    "../../../../documentation/source/images/examples/street_busy.jpg",
    "https://cdn-attachments.timesofmalta.com/cc1eceadde40d2940bc5dd20692901371622153217-1301777007-4d978a6f-620x348.jpg",
]

predictions = model.predict(IMAGES)
predictions.show()
predictions.save(output_folder="")  # Save in working directory

Exception

    shift_x = torch.arange(end=w, dtype=dtype) + self.grid_cell_offset
RuntimeError: "arange_cpu" not implemented for 'Half'

pytorch == '1.12.0+cu102'

Solutions

1. remove dtype
2. check pytorch version and use dtype only if user has the right version. I am not sure what's the best way to find which version supports this, except by trying them one by one.

Should we go for 1. or do we still want to add dtype when possible ? What was the motivation to add it ? @BloodAxe

[BloodAxe]

Dtype is needed since otherwise it causes fp64 types appear in onnx graph.
But I'm not sure where fp16 issue is coming from. I will take a look, thanks for raising this issue

[Pbatch]

Is it possible to add just the device to torch.arange instead?

I.e.

shift_x = torch.arange(end=w, device=device, dtype=dtype) + self.grid_cell_offset

Creating a torch.float16 tensor is fine if it's made on the GPU.

[Louis-Dupont]

LGTM

[BloodAxe]

LGTM