code: def wait(
      batch: Batch,
      prev_stream: AbstractStream,
      next_stream: AbstractStream,
  ) -> None:
      batch[:] = Wait.apply(prev_stream, next_stream, *batch)
      # Gradients are only supported for float Tensors.
      batch[:] = tuple(
          [x if x.is_floating_point() else x.detach() for x in batch]
      )
  
imports: from queue import Queue
  from types import TracebackType
  from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Type, Union, cast
  import torch
  from torch import Tensor, nn
  from torch.autograd.profiler import record_function
  from .checkpoint import Checkpointing
  from .copy import Copy, Wait
  from .dependency import fork, join
  from .microbatch import Batch
  from .skip.layout import SkipLayout
  from .skip.tracker import SkipTrackerThroughPotals, use_skip_tracker
  from .stream import AbstractStream, current_stream, use_device
  from .worker import Task, create_workers, join_workers
label_desc: Use image node_id 3 for calling a global function with example usage: wait(batch, prev_stream, next_stream) without return types