code: def dropout_outshape(module_masks, mask):
      if module_masks.output_mask is None:
          module_masks.set_output_mask(mask)
          module_masks.set_input_mask(mask)
          return module_masks.input_mask
      # if alreay visited
      assert all(
          module_masks.output_mask.mask_index[1] == mask.mask_index[1]
      )
  
      return module_masks.output_mask
  
imports: import logging
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: dropout_outshape(module_masks, mask) and returns: module_masks, module_masks