Thank you.
Could you check if you observe the “Grad is None” behaviour if you do the following:
src_batch = torch.randn(5, 128, 128)
tar_batch = torch.randn(5, 128, 128)

model = Net(8)
run_model(model, src_batch, tar_batch)