code: def test_reset_peak_memory_stats(
      prev_max_memory: int, allocated_memory: int
  ) -> None:
      """Tests resetting of peak memory stats.
  
      Notes: Only the relative and not the absolute sizes of prev_max_memory and
      allocated_memory are relevant.
  
      Args:
          prev_max_memory: current maximum memory stat to simulate
          allocated_memory: current allocated memory to simulate
      """
      device = torch.device("cuda:0")
  
      # keep x, delete y
      x = get_n_byte_tensor(allocated_memory, device=device)
      y = get_n_byte_tensor(
          prev_max_memory - allocated_memory, device=device
      )
      del y
  
      # get the true allocated memory (CUDA memory is allocated in blocks)
      prev_max_memory = torch.cuda.max_memory_allocated(device)
      allocated_memory = torch.cuda.memory_allocated(device)
      assert prev_max_memory >= allocated_memory
      assert reset_peak_memory_stats(device) == (
          prev_max_memory,
          allocated_memory,
      )
  
      # clean up
      del x
      torch.cuda.reset_peak_memory_stats(device)
      assert torch.cuda.max_memory_allocated(device) == 0
      assert torch.cuda.memory_allocated(device) == 0
  
file_path: opacus/benchmarks/tests/test_utils.py
imports: import os
  import pickle
  import shutil
  from typing import Any, Dict, List, Tuple
  import pytest
  import torch
  from helpers import get_n_byte_tensor, skipifnocuda
  from benchmarks.utils import get_layer_set, get_path, reset_peak_memory_stats, save_results
label_desc: Use image node_id 2 for calling a global function with example usage: test_reset_peak_memory_stats(prev_max_memory, allocated_memory) without return types