It’s the processing. Namely the last line that feeds the input to the model. I’ve gotten as far as the parallel_apply which does not return.
py(124)parallel_apply()
-> return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
(Pdb) 
	In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

Even going deeper, I can trace it down to thread.start() inside parallel_apply.
And further down,
> /usr/lib/python3.6/threading.py(828)start()
-> def start(self):
(Pdb) n
> /usr/lib/python3.6/threading.py(838)start()
-> if not self._initialized:
(Pdb) 
> /usr/lib/python3.6/threading.py(841)start()
-> if self._started.is_set():
(Pdb) 
> /usr/lib/python3.6/threading.py(843)start()
-> with _active_limbo_lock:
(Pdb) 
> /usr/lib/python3.6/threading.py(844)start()
-> _limbo[self] = self
(Pdb) 
> /usr/lib/python3.6/threading.py(845)start()
-> try:
(Pdb) 
> /usr/lib/python3.6/threading.py(846)start()
-> _start_new_thread(self._bootstrap, ())
(Pdb) s
> /usr/lib/python3.6/threading.py(851)start()
-> self._started.wait()
(Pdb) 	In Model: input size torch.Size([15, 5]) output size torch.Size([15, 2])

At this point, it hangs wtihout stepping in. I’m guessing that this is mapped to pthread call. I can fire up gdb and see where it leads. However it would be counterproductive if it is a known issue and/or I’m just doing something silly.