File size: 485 Bytes
9b19c29 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 |
from typing import overload
import torch
@overload
def to_optional_float(x: torch.Tensor) -> float:
...
@overload
def to_optional_float(x: float) -> float:
...
@overload
def to_optional_float(x: None) -> None:
...
def to_optional_float(x: torch.Tensor | float | None) -> float | None:
"""For the common case where one needs to extract a float from a scalar Tensor, which may be None."""
if isinstance(x, torch.Tensor):
return x.item()
return x
|