Pytorch mdoule translation to tensorflow

I was trying to translate a PyTorch code (a single module) to TensorFlow but faced some obstacles. Below are the both code and trails. How should I do this?

In pytorch, it has

import torch.nn as nn
import torch

class Blcoks(nn.Module):
    def __init__(self):

    def forward(self, local_feat, global_feat):
        global_feat_norm = torch.norm(global_feat, p=2, dim=1) 

        projection = torch.bmm(global_feat.unsqueeze(1), torch.flatten(local_feat, start_dim=2))
        projection = torch.bmm(global_feat.unsqueeze(2), projection).view(local_feat.size())
        projection = projection / (global_feat_norm * global_feat_norm).view(-1, 1, 1, 1)

        orthogonal_comp = local_feat - projection
        global_feat = global_feat.unsqueeze(-1).unsqueeze(-1)
                          orthogonal_comp], dim=1)

local_  = torch.ones(1, 20, 32, 32)
global_ = torch.ones(1, 20)
x = OrthogonalFusion()
x(local_, global_).shape
torch.Size([1, 40, 32, 32])

Above is the working code. Now, I’ve tried to translate into tf but fails.

Major issue,

Is it possible to translate? I’m mostly used to pytorch but now I’m thinking to switch to tensorflwo but translation is not going smooth.

If you have problem with flatten you could try to see: