Fix 38_ConvTranspose3d_AvgPool_Clamp_Softmax_Multiply.py

Author: nataliakokoromyti

KernelBench/level2/38_ConvTranspose3d_AvgPool_Clamp_Softmax_Multiply.py

@@ -3,14 +3,16 @@
 
 class Model(nn.Module):
     """
-    Model that performs a 3D transposed convolution, average pooling, clamping, softmax, and multiplication.
+    Model that performs average pooling, 3D transposed convolution, clamping,
+    spatial softmax, and multiplication by a learnable scale.
     """
     def __init__(self, in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, clamp_min, clamp_max):
         super(Model, self).__init__()
-        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
         self.avg_pool = nn.AvgPool3d(pool_kernel_size)
+        self.conv_transpose = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding, output_padding=output_padding)
         self.clamp_min = clamp_min
         self.clamp_max = clamp_max
+        self.scale = nn.Parameter(torch.ones(1, out_channels, 1, 1, 1))
 
     def forward(self, x):
         """
@@ -20,11 +22,14 @@ def forward(self, x):
         Returns:
             torch.Tensor: Output tensor of shape (batch_size, out_channels, depth, height, width).
         """
-        x = self.conv_transpose(x)
         x = self.avg_pool(x)
+        x = self.conv_transpose(x)
         x = torch.clamp(x, self.clamp_min, self.clamp_max)
-        x = torch.softmax(x, dim=1)
-        x = x * 2
+        b, c, d, h, w = x.shape
+        x = x.view(b, c, -1)                     # flatten spatial dims
+        x = torch.softmax(x, dim=2)
+        x = x.view(b, c, d, h, w)
+        x = x * self.scale
         return x
 
 batch_size = 16
@@ -43,4 +48,4 @@ def get_inputs():
     return [torch.randn(batch_size, in_channels, depth, height, width)]
 
 def get_init_inputs():
-    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, clamp_min, clamp_max]
+    return [in_channels, out_channels, kernel_size, stride, padding, output_padding, pool_kernel_size, clamp_min, clamp_max]