Non-tensor module outputs

[vieting]

I have a case where the output of a module is not a tensor. A simple example could be:

class DummyModule(torch.nn.Module):
  def forward(self, y):
    return {"a": 1, "b": y + 1}

I'm not exactly sure about all the logic in CallEntry, but just focusing on the Tensor outputs seems to work here.