Update ResNet.py

toummHus · Aug 16, 2024 · 84cb63f · 84cb63f
1 parent d2cd41c
commit 84cb63f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/utils/ResNet.py b/utils/ResNet.py
@@ -245,7 +245,8 @@ def forward(self, x):
         out = self.proj_out(out)
         return x+out
 
-
+def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
+    return AttnBlock(in_channels)
 
 def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
     assert attn_type in ["vanilla", "vanilla-xformers", "memory-efficient-cross-attn", "linear", "none"], f'attn_type {attn_type} unknown'