Hi, I have a problem compiling a model for hailo8. It is a larger model and for one layer I defined a Correlation Block, that includes several transpose, view and a torch.bmm Operation. So far, I learned that torch.bmm cannot be used on hailo8. And I learned that my tensor reshaping operations are difficult. However, I hope there is a possibility to change this class such that I can convert it for hailo8. Do you have any suggestions? The input shapes are [1, 768, 24, 24] and [1,768, 8,8] (not too large…) class CorrelationBlock(nn.Module):
def init(self):
super().init()
self.ReLU = nn.ReLU(inplace=True)
def forward(self, x, y):
b_x, c_x, h_x, w_x = x.size()
b_y, c_y, h_y, w_y = y.size()
feature_x = x.transpose(2, 3).contiguous().view(b_x, c_x, h_x * w_x)
feature_y = y.view(b_y, c_y, h_y * w_y).transpose(1, 2)
feature_mul = torch.bmm(feature_y, feature_x)
feature_mul_trans = feature_mul.view(b_x, h_y, w_y, h_x * w_x).transpose(2, 3).transpose(1, 2)
return self.ReLU(feature_mul_trans)