code: def __init__(
      self,
      n_steps: int,
      n_heads: int,
      n_feat: int,
      emb_size: int,
      emb_dim: int,
  ):
      super().__init__()
      d_model = n_feat * n_heads
      self.layer_norm1 = torch.nn.LayerNorm(d_model)
      self.attention1 = AlphaMultiHeadAttention(d_model, d_model)
      self.drop1 = torch.nn.Dropout()
      self.layer_norm2 = torch.nn.LayerNorm(d_model)
      self.attention2 = AlphaMultiHeadAttention(d_model, emb_dim)
      self.drop2 = torch.nn.Dropout()
  
imports: import math
  import torch
  import torch.nn.functional
  from open_alpha_tensor.core.modules.attention import AlphaMultiHeadAttention
label_desc: Use image node_id 1 to create a new PolicyHeadDoubleAttention object from inherited base classes: torch.nn with example: obj = PolicyHeadDoubleAttention(n_steps, n_heads, n_feat, emb_size, emb_dim)