From 7dd2916abdaccc0b5c9b012dd505eaec4cd11338 Mon Sep 17 00:00:00 2001
From: Junming Chen <72211694+Leoooo333@users.noreply.github.com>
Date: Tue, 15 Aug 2023 11:29:17 -0400
Subject: [PATCH] Update torch.py

---
 equilib/equi2pers/torch.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/equilib/equi2pers/torch.py b/equilib/equi2pers/torch.py
index 1de02662..5ebf0923 100644
--- 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)
     )