Remove DeMemWM attention debugger trap

algorithms/dememwm/models/attention.py

@@ -94,17 +94,6 @@ class TemporalAxialAttention(nn.Module):
     def forward(self, x: torch.Tensor, frame_memory_segments=None, frame_memory_masks=None):
         B, T, H, W, D = x.shape
 
-        # if T>=9:
-        #     try:
-        #         # x = torch.cat([x[:,:-1],x[:,16-T:17-T],x[:,-1:]], dim=1)
-        #         x = torch.cat([x[:,16-T:17-T],x], dim=1)
-        #     except:
-        #         import pdb;pdb.set_trace()
-        #     print("="*50)
-        #     print(x.shape)
-
-        B, T, H, W, D = x.shape
-
         q, k, v = self.to_qkv(x).chunk(3, dim=-1)
 
         if self.use_domain_adapter:
@@ -145,24 +134,13 @@ class TemporalAxialAttention(nn.Module):
         else:
             attn_bias = None
 
-        try:
-            x = F.scaled_dot_product_attention(query=q, key=k, value=v, attn_mask=attn_bias)
-        except:
-            import pdb;pdb.set_trace()
+        x = F.scaled_dot_product_attention(query=q, key=k, value=v, attn_mask=attn_bias)
 
         x = rearrange(x, "(B H W) h T d -> B T H W (h d)", B=B, H=H, W=W)
         x = x.to(q.dtype)
 
         # linear proj
         x = self.to_out(x)
-
-        # if T>=10:
-        #     try:
-        #         # x = torch.cat([x[:,:-2],x[:,-1:]], dim=1)
-        #         x = x[:,1:]
-        #     except:
-        #         import pdb;pdb.set_trace()
-        #     print(x.shape)
         return x
 
 class SpatialAxialAttention(nn.Module):