Update torch.py

haruishi43 · Aug 15, 2023 · 7dd2916 · 7dd2916
1 parent 49c92c6
commit 7dd2916
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/equilib/equi2pers/torch.py b/equilib/equi2pers/torch.py
@@ -117,6 +117,7 @@ def run(
     z_down: bool,
     mode: str,
     backend: str = "native",
+    clip_output: bool = True,
 ) -> torch.Tensor:
     """Run Equi2Pers
 
@@ -242,7 +243,7 @@ def run(
 
     out = (
         out.type(equi_dtype)
-        if equi_dtype == torch.uint8
+        if equi_dtype == torch.uint8 or not clip_output
         else torch.clip(out, 0.0, 1.0)
     )