text: class STPN(nn.Module):
      def __init__(self, height_feat_size=13):
          super(STPN, self).__init__()
          self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
          self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
          self.bn_pre_1 = nn.BatchNorm2d(32)
          self.bn_pre_2 = nn.BatchNorm2d(32)
  
  #        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
  #        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
          self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
  
          self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
          self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
          self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
          self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
          self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
  
          self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
          self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
  
          self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
          self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
  
          self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
          self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
  
          self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
          self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
  
          self.bn1_1 = nn.BatchNorm2d(64)
          self.bn1_2 = nn.BatchNorm2d(64)
  
          self.bn2_1 = nn.BatchNorm2d(128)
          self.bn2_2 = nn.BatchNorm2d(128)
  
          self.bn3_1 = nn.BatchNorm2d(256)
          self.bn3_2 = nn.BatchNorm2d(256)
  
          self.bn4_1 = nn.BatchNorm2d(512)
          self.bn4_2 = nn.BatchNorm2d(512)
  
          self.bn5_1 = nn.BatchNorm2d(256)
          self.bn5_2 = nn.BatchNorm2d(256)
  
          self.bn6_1 = nn.BatchNorm2d(128)
          self.bn6_2 = nn.BatchNorm2d(128)
  
          self.bn7_1 = nn.BatchNorm2d(64)
          self.bn7_2 = nn.BatchNorm2d(64)
  
          self.bn8_1 = nn.BatchNorm2d(32)
          self.bn8_2 = nn.BatchNorm2d(32)
  
      def forward(self, x):
          batch, seq, z, h, w = x.size()
  
          x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
          x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
          x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
  
          # -------------------------------- Encoder Path --------------------------------
          # -- STC block 1
          x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
          x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_1 = self.conv3d_1(x_1)
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()  # (batch * seq, c, h, w)
  
          # -- STC block 2
          x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
          x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3)).contiguous()  # (batch, seq, c, h, w)
          x_2 = self.conv3d_2(x_2)
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()  # (batch * seq, c, h, w), seq = 1
  
          # -- STC block 3
          x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
          x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))
  
          # -- STC block 4
          x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
          x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))
  
  
          # -------------------------------- Decoder Path --------------------------------
          x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
          x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))
  
          x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
          x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
          x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()
  
          x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
          x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))
  
          x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
          x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
          x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()
  
          x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
          x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
  
          x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = F.adaptive_max_pool3d(x, (1, None, None))
          x = x.permute(0, 2, 1, 3, 4).contiguous()
          x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()
  
          x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
          res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
  
          return res_x