The annotation of the parameter device of PyTorchPredictor is incomplete

[wu462]

Description

The annotation of the parameter device of PyTorchPredictor is incomplete, it should be Union[str, torch.device]

To Reproduce

class PyTorchPredictor(RepresentablePredictor):
    @validated()
    def __init__(
        self,
        input_names: List[str],
        prediction_net: nn.Module,
        batch_size: int,
        prediction_length: int,
        input_transform: Transformation,
        forecast_generator: ForecastGenerator = SampleForecastGenerator(),
        output_transform: Optional[OutputTransform] = None,
        lead_time: int = 0,
        device: str = "auto",
    ) -> None:

Error message or code output

pydantic.v1.error_wrappers.ValidationError: 1 validation error for PyTorchPredictorModel
device
  str type expected (type=type_error.str)

[lostella]

For convenience, this is the line in question

gluonts/src/gluonts/torch/model/predictor.py

Line 53 in 4c8eefd

device: str = "auto",

[cooling-0]

Thank youuuuuuuuuuuuuuuuuuu!
I had the same issue.
But thanks to your question I fixed it.