code: def exception_no_hang(pipe_class):
      # In v0.0.2, once a failed partition receives a normal message
      # (non-closing) for the next micro-batch, a hang occured. The reason was
      # that a failed partition didn't call in_queue.task_done() on a normal
      # message. So the former partition was blocked at out_queue.join() for the
      # next of next micro-batch.
      class ExpectedException(Exception):
          pass
  
      class Pass(nn.Module):
          def forward(self, x):
              return x
  
      class Raise(nn.Module):
          def forward(self, x):
              raise ExpectedException()
  
      model = nn.Sequential(Pass(), Pass(), Raise())
      model = pipe_class(
          model, [1, 1, 1], worker_map=get_worker_map(), chunks=3
      )
      model.eval()
  
      if model.group.rank() == 2:
          with pytest.raises(ExpectedException):
              model(torch.rand(3))
      else:
          model(torch.rand(3))
  
      torch.distributed.barrier()
  
file_path: fairscale/tests/nn/pipe_process/test_bugs.py
imports: import pytest
  import torch
  from torch import nn
  import torch.nn.functional
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.nn.pipe import AsyncPipe
label_desc: Use image node_id 2 for calling a global function with example usage: exception_no_hang(pipe_class) and returns: x