address comment

pyro-ppl · Mar 4, 2023 · e74831c · e74831c
1 parent 5c5f1bb
commit e74831c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyro/distributions/torch_transform.py b/pyro/distributions/torch_transform.py
@@ -35,6 +35,6 @@ def __hash__(self):
         return super(torch.nn.Module, self).__hash__()
 
     def with_cache(self, cache_size=1):
-        if cache_size == 0:
+        if cache_size == self._cache_size:
             return self
         return ComposeTransformModule(self.parts, cache_size=cache_size)