addressing Torchscript compatibility

[SergeySandler]

Making the code torch.jit.script() friendly; torch.jit.trace() is also supported.
See issues/83.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[SergeySandler]

If there is a use case to return unrefined_occlusion, unrefined_tracks and unrefined_expected_dist then the code should be modified as

    out = (torch.mean(torch.stack(trajectories.occlusion[p::p]), dim=0),       # occlusion
           torch.mean(torch.stack(trajectories.tracks[p::p]), dim=0),          # tracks
           torch.mean(torch.stack(trajectories.expected_dist[p::p]), dim=0),   # expected_dist
           trajectories.occlusion[:-1],                                        # unrefined_occlusion
           trajectories.tracks[:-1],                                           # unrefined_tracks
           trajectories.expected_dist[:-1]                                     # unrefined_expected_dist        
          )

having -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: replaced with -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: in the forward annotation.