Trying attention method

2026-04-26 22:32:13 +01:00
parent a2363d882f
commit 80a314b9e9
1 changed files with 2 additions and 1 deletions
@@ -139,7 +139,8 @@ class ShepherdAttentionExtractor(BaseFeaturesExtractor):
            dropout=0.0, batch_first=True,
        )
        self.transformer = nn.TransformerEncoder(encoder_layer,
-                                                 num_layers=n_layers)
+                                                 num_layers=n_layers,
+                                                 enable_nested_tensor=False)

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        B = obs.shape[0]