code: def __init__(
      self,
      mask_type: str,
      in_channels: int,
      out_channels: int,
      kernel_size: int,
      stride: int = 1,
      padding: Union[str, int] = 0,
      dilation: int = 1,
      groups: int = 1,
      bias: bool = True,
      padding_mode: str = "zeros",
      device=None,
      dtype=None,
  ) -> None:
      super().__init__(
          in_channels,
          out_channels,
          kernel_size,
          stride,
          padding,
          dilation,
          groups,
          bias,
          padding_mode,
          device,
          dtype,
      )
  
      self.register_buffer("mask", torch.ones_like(self.weight))
  
      _, _, kH, kW = self.weight.shape
  
      if mask_type == "A":
          self.mask[:, :, kH // 2, kW // 2 :] = 0
          self.mask[:, :, kH // 2 + 1 :] = 0
  
      else:
          self.mask[:, :, kH // 2, kW // 2 + 1 :] = 0
          self.mask[:, :, kH // 2 + 1 :] = 0
  
file_path: benchmark_VAE/src/pythae/models/normalizing_flows/pixelcnn/utils.py
imports: from typing import Union
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new MaskedConv2d object from inherited base classes: nn with example: obj = MaskedConv2d(mask_type, in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode, device, dtype)