From e95e514176716c27db2284043d45412910d1a550 Mon Sep 17 00:00:00 2001
From: Sarah Tran <sarahtran@meta.com>
Date: Fri, 13 Jun 2025 09:44:31 -0700
Subject: [PATCH] Pin click version to unblock pyre in OSS tests (#1585)

Summary:
Pull Request resolved: https://github.com/pytorch/captum/pull/1585

Rollback Plan:

Differential Revision: D76605474
---
 captum/module/stochastic_gates_base.py | 2 ++
 setup.py                               | 1 +
 2 files changed, 3 insertions(+)

diff --git a/captum/module/stochastic_gates_base.py b/captum/module/stochastic_gates_base.py
index b34a4d5f4d..7c9d752c08 100644
--- a/captum/module/stochastic_gates_base.py
+++ b/captum/module/stochastic_gates_base.py
@@ -30,6 +30,8 @@ class StochasticGatesBase(Module, ABC):
     extend this class and implement the distribution specific functions.
     """
 
+    mask: Optional[Tensor]
+
     def __init__(
         self,
         n_gates: int,
diff --git a/setup.py b/setup.py
index bb11265897..f6c62d8a0d 100644
--- a/setup.py
+++ b/setup.py
@@ -81,6 +81,7 @@ def report(*args):
         "ufmt",
         "scikit-learn",
         "annoy",
+        "click<8.2.0",
     ]
 )