Move EfficientAttentionConfig to module level

sparktts/modules/speaker/perceiver_encoder.py

@@ -45,6 +45,21 @@ def once(fn):
 
 print_once = once(print)
 
+# Define config class at module level
+class EfficientAttentionConfig:
+    def __init__(self, enable_flash, enable_math, enable_mem_efficient):
+        self.enable_flash = enable_flash
+        self.enable_math = enable_math
+        self.enable_mem_efficient = enable_mem_efficient
+
+    def _asdict(self):
+        return {
+            'enable_flash': self.enable_flash,
+            'enable_math': self.enable_math,
+            'enable_mem_efficient': self.enable_mem_efficient
+        }
+
+
 # main class
 
 
@@ -62,20 +77,6 @@ class Attend(nn.Module):
             use_flash and version.parse(torch.__version__) < version.parse("2.0.0")
         ), "in order to use flash attention, you must be using pytorch 2.0 or above"
 
-        # Define config as a regular class instead of namedtuple
-        class EfficientAttentionConfig:
-            def __init__(self, enable_flash, enable_math, enable_mem_efficient):
-                self.enable_flash = enable_flash
-                self.enable_math = enable_math
-                self.enable_mem_efficient = enable_mem_efficient
-
-            def _asdict(self):
-                return {
-                    'enable_flash': self.enable_flash,
-                    'enable_math': self.enable_math,
-                    'enable_mem_efficient': self.enable_mem_efficient
-                }
-
         self.cpu_config = EfficientAttentionConfig(True, True, True)
         self.cuda_config = None
 
@@ -88,7 +89,7 @@ class Attend(nn.Module):
             print_once(
                 "A100 GPU detected, using flash attention if input tensor is on cuda"
             )
-            self.cuda_config = self.config(True, False, False)
+            self.cuda_config = EfficientAttentionConfig(True, False, False)
         else:
             print_once(
                 "Non-A100 GPU detected, using math or mem efficient attention if input tensor is on cuda"