changed type hinting

atong01 · Aug 9, 2024 · 41cd14c · 41cd14c
1 parent 24de1a2
commit 41cd14c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchcfm/optimal_transport.py b/torchcfm/optimal_transport.py
@@ -1,7 +1,7 @@
 import math
 import warnings
 from functools import partial
-from typing import Optional
+from typing import Optional, Union
 
 import numpy as np
 import ot as pot
@@ -18,7 +18,7 @@ def __init__(
         reg: float = 0.05,
         reg_m: float = 1.0,
         normalize_cost: bool = False,
-        num_threads: int | str = 1,
+        num_threads: Union[int, str] = 1,
         warn: bool = True,
     ) -> None:
         """Initialize the OTPlanSampler class.