class_docstr: An implementation of the spatial-temporal attention block, with spatial attention and temporal attention
  followed by gated fusion. For details see this paper: `"GMAN: A Graph Multi-Attention Network for Traffic Prediction."
  <https://arxiv.org/pdf/1911.08415.pdf>`_
  
  Args:
      K (int) : Number of attention heads.
      d (int) : Dimension of each attention head outputs.
      bn_decay (float): Batch normalization momentum.
      mask (bool): Whether to mask attention score in temporal attention.
docstr: Making a forward pass of the spatial-temporal attention block.
  
  Arg types:
      * **X** (PyTorch Float Tensor) - Input sequence, with shape (batch_size, num_step, num_nodes, K*d).
      * **STE** (Pytorch Float Tensor) - Spatial-temporal embedding, with shape (batch_size, num_step, num_nodes, K*d).
  
  Return types:
      * **X** (PyTorch Float Tensor) - Attention scores, with shape (batch_size, num_step, num_nodes, K*d).
docstr_tok: ["Making","a","forward","pass","of","the","spatial-temporal","attention","block",".","Arg","types",":","*","*","*","X","*","*","(","PyTorch","Float","Tensor",")","-","Input","sequence",",","with","shape","(","batch_size",",","num_step",",","num_nodes",",","K","*","d",")",".","*","*","*","STE","*","*","(","Pytorch","Float","Tensor",")","-","Spatial-temporal","embedding",",","with","shape","(","batch_size",",","num_step",",","num_nodes",",","K","*","d",")",".","Return","types",":","*","*","*","X","*","*","(","PyTorch","Float","Tensor",")","-","Attention","scores",",","with","shape","(","batch_size",",","num_step",",","num_nodes",",","K","*","d",")","."]
code: def forward(
      self, X: torch.FloatTensor, STE: torch.FloatTensor
  ) -> torch.FloatTensor:
      """
      Making a forward pass of the spatial-temporal attention block.
  
      Arg types:
          * **X** (PyTorch Float Tensor) - Input sequence, with shape (batch_size, num_step, num_nodes, K*d).
          * **STE** (Pytorch Float Tensor) - Spatial-temporal embedding, with shape (batch_size, num_step, num_nodes, K*d).
  
      Return types:
          * **X** (PyTorch Float Tensor) - Attention scores, with shape (batch_size, num_step, num_nodes, K*d).
      """
      HS = self._spatial_attention(X, STE)
      HT = self._temporal_attention(X, STE)
      H = self._gated_fusion(HS, HT)
      del HS, HT
      X = torch.add(X, H)
      return X
  
imports: import math
  from typing import Union, Callable, Optional
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the SpatioTemporalAttention obj's underlying member method code with example usage: obj.forward(X, STE) and returns: X