Indexed assignment not getting recorded in traced model

1

Hi,

I am creating a zeros tensor and then assigning other tensors to it by indexing into the original tensor like so-

def f(b): 
            num = b.size()[0]    
            a = b.new_zeros(num, 5) 
            a[:, :4] = b 
            return a

traced_func = torch.jit.trace(f, torch.randn(20, 4))
print (traced_func.code)
def forward(self,
    b: Tensor) -> Tensor:
  a = torch.zeros([20, 5], dtype=6, layout=0, device=torch.device("cpu"))
  return a

As you can see the indexed assignment of tensor b to tensor a is not getting recorded in the traced graph