code: def test_exception_no_hang():
      # 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(
          model, [1, 1, 1], devices=["cpu", "cpu", "cpu"], chunks=3
      )
  
      with pytest.raises(ExpectedException):
          model(torch.rand(3))
  
file_path: fairscale/tests/nn/pipe/test_bugs.py
imports: import pytest
  import torch
  from torch import nn
  import torch.nn.functional
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu
  from fairscale.nn.pipe import Pipe
label_desc: Use image node_id 2 for calling a global function with example usage: test_exception_no_hang() and returns: x