[FIX FR NOW]

kyegomez · Jan 24, 2024 · 1018e07 · 1018e07
1 parent 10ffe02
commit 1018e07
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/swarms_torch/drone_swarm.py b/swarms_torch/drone_swarm.py
@@ -309,9 +309,12 @@ def forward(
         final_neighborhood_embedding = self.neighbor_mlp(obs_neighbors)
         return final_neighborhood_embedding
 
+@dataclass
+class SwarmMultiHeadAttentionEncoder(nn.Module):
+    dim: int
 
 @dataclass
-class SwarmMultiHeadAttention(SwarmMultiHeadAttentionEncoder):
+class QuadSingleHeadAttentionEncoderSim2Real(SwarmMultiHeadAttentionEncoder):
     obs_space: int
     quads_obs_repr: Any
     neighbor_hidden_size: int