code: def _hook_memory():
      if torch.cuda.is_available():
          torch.cuda.synchronize()
          used_memory = torch.cuda.max_memory_allocated()
      else:
          used_memory = np.nan
      return used_memory
  
file_path: pytorch-lightning/tests/parity_pytorch/test_basic_parity.py
imports: import os
  import numpy
  import pytest
  import torch
  from lightning.pytorch import LightningModule, Trainer, seed_everything
  from tests_pytorch.helpers.advanced_models import ParityModuleMNIST, ParityModuleRNN
  from parity_pytorch.measure import measure_loops
  from parity_pytorch.models import ParityModuleCIFAR
label_desc: Use image node_id 4 for calling a global function with example usage: _hook_memory() and returns: used_memory