Multiplying 2 tensors properly

Hi team,

i need to mutiply to images and pass to a conv layer

shape of image is = torch.Size([2, 256, 16, 16])
style = torch.Size([2, 1, 1, 256])

How can i mutiply both the shape and style with proper output to produce

Code i am using to solve -

self.style_scale_transform = nn.Linear(256,256)
self.style_shift_transform = nn.Linear(256,256)
style_scale = self.style_scale_transform(style).unsqueeze(4).unsqueeze(5)
style_bias = self.style_shift_transform(style).unsqueeze(4).unsqueeze(5)
out = x * style_scale + style_bias

Please let me know the shape of resulting output.

I guess you can do this,

# (2, 256, 16, 16) -> (2, 256, 256) -> (2, 256, 256)
image = image.flatten(start_dim=2).transpose(2, 1)
# (2, 1, 1, 256) -> (2, 1, 256) -> (2, 256, 1)
style = style.squeeze(dim=1).transpose(2, 1)

# (2, 256, 256) x (2, 256, 1) = (2, 256, 1)
result = torch.bmm(image, style)

Hope it helps you

but wont this disturb the image dimnesions and final output