diff --git a/grassmann_tensor/tensor.py b/grassmann_tensor/tensor.py index ab2eedc..c1bd774 100644 --- a/grassmann_tensor/tensor.py +++ b/grassmann_tensor/tensor.py @@ -68,19 +68,18 @@ def to(self, whatever: torch.device | torch.dtype | str | None = None, *, device """ Copy the tensor to a specified device or copy it to a specified data type. """ - if whatever is None: - pass - elif isinstance(whatever, torch.device): - assert device is None, "Duplicate device specification." - device = whatever - elif isinstance(whatever, torch.dtype): - assert dtype is None, "Duplicate dtype specification." - dtype = whatever - elif isinstance(whatever, str): - assert device is None, "Duplicate device specification." - device = torch.device(whatever) - else: - raise TypeError(f"Unsupported type for 'to': {type(whatever)}. Expected torch.device, torch.dtype, or str.") + match whatever: + case torch.device(): + assert device is None, "Duplicate device specification." + device = whatever + case torch.dtype(): + assert dtype is None, "Duplicate dtype specification." + dtype = whatever + case str(): + assert device is None, "Duplicate device specification." + device = torch.device(whatever) + case None: + pass match (device, dtype): case (None, None): return self diff --git a/tests/conversion_test.py b/tests/conversion_test.py index b9f2816..3454319 100644 --- a/tests/conversion_test.py +++ b/tests/conversion_test.py @@ -42,11 +42,6 @@ def test_conversion( y = x.to(*args, **kwargs) -def test_conversion_invalid_type(x: GrassmannTensor) -> None: - with pytest.raises(TypeError): - x.to(2333) # type: ignore[arg-type] - - def test_conversion_duplicated_value(x: GrassmannTensor) -> None: with pytest.raises(AssertionError): x.to(torch.device("cpu"), device=torch.device("cpu"))