property for sigma int/float

atong01 · Nov 13, 2023 · b12f8fa · b12f8fa
1 parent ed6b607
commit b12f8fa
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/torchcfm/conditional_flow_matching.py b/torchcfm/conditional_flow_matching.py
@@ -54,7 +54,16 @@ def __init__(self, sigma: float = 0.0):
         ----------
         sigma : float
         """
-        self.sigma = sigma
+        self._sigma = sigma
+
+    @property
+    def sigma(self):
+        if isinstance(self._sigma, float):
+            return self._sigma
+        elif isinstance(self._sigma, int):
+            return float(self._sigma)
+        else:
+            raise ValueError("Sigma must be a float or int.")
 
     def compute_mu_t(self, x0, x1, t):
         """