Update src/attentionhacked_garmnet.py

src/attentionhacked_garmnet.py

@@ -55,6 +55,7 @@ def _chunked_feed_forward(
 class GatedSelfAttentionDense(nn.Module):
     r"""
     A gated self-attention dense layer that combines visual features and object features.
+
     Parameters:
         query_dim (`int`): The number of channels in the query.
         context_dim (`int`): The number of channels in the context.
@@ -96,6 +97,7 @@ class GatedSelfAttentionDense(nn.Module):
 class BasicTransformerBlock(nn.Module):
     r"""
     A basic Transformer block.
+
     Parameters:
         dim (`int`): The number of channels in the input and output.
         num_attention_heads (`int`): The number of heads to use for multi-head attention.
@@ -408,6 +410,7 @@ class BasicTransformerBlock(nn.Module):
 class TemporalBasicTransformerBlock(nn.Module):
     r"""
     A basic Transformer block for video like data.
+
     Parameters:
         dim (`int`): The number of channels in the input and output.
         time_mix_inner_dim (`int`): The number of channels for temporal attention.
@@ -609,6 +612,7 @@ class SkipFFTransformerBlock(nn.Module):
 class FeedForward(nn.Module):
     r"""
     A feed-forward layer.
+
     Parameters:
         dim (`int`): The number of channels in the input.
         dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
@@ -663,4 +667,4 @@ class FeedForward(nn.Module):
                 hidden_states = module(hidden_states, scale)
             else:
                 hidden_states = module(hidden_states)
-        return hidden_states
+        return hidden_states