p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/nn.py (python)
        z_BTNM = self.spatial_pos_enc(x_BTNM)
        z_BTNM = self.spatial_norm(z_BTNM)
        z_BTNM = self.spatial_attention(z_BTNM)
        x_BTNM = x_BTNM + z_BTNM

        # --- Temporal attention ---
        x_BNTM = x_BTNM.swapaxes(1, 2)
        z_BNTM = self.temporal_pos_enc(x_BNTM)
        z_BNTM = self.temporal_norm(z_BNTM)
        z_BNTM = self.temporal_attention(z_BNTM)
        x_BNTM = x_BNTM + z_BNTM
        x_BTNM = x_BNTM.swapaxes(1, 2)