Change WarningCache to subclass set (#7995)

Lightning-AI · Jun 16, 2021 · 2134216 · 2134216
1 parent 4ffba60
commit 2134216
Showing 1 changed file with 3 additions and 9 deletions.
diff --git a/pytorch_lightning/utilities/warnings.py b/pytorch_lightning/utilities/warnings.py
@@ -14,15 +14,9 @@
 from pytorch_lightning.utilities.distributed import rank_zero_warn
 
 
-class WarningCache:
-
-    def __init__(self):
-        self.warnings = set()
+class WarningCache(set):
 
     def warn(self, m, *args, **kwargs):
-        if m not in self.warnings:
-            self.warnings.add(m)
+        if m not in self:
+            self.add(m)
             rank_zero_warn(m, *args, **kwargs)
-
-    def clear(self):
-        self.warnings.clear()