code: def get_loss_and_gnorm(model, input):
      """Helper to run a forward/backward pass and return results in a dict."""
      ret = {}
  
      ret["mem_0"] = get_cuda_mem_allocated()
      ret["mem_peak"] = 0
      if ret["mem_0"] > 0:
          torch.cuda.reset_peak_memory_stats()
  
      model.zero_grad()
      loss = model(input).sum()
      ret["mem_after_fwd"] = get_cuda_mem_allocated()
  
      loss.backward()
      ret["mem_after_bwd"] = get_cuda_mem_allocated()
  
      gnorm = torch.norm(
          torch.stack(
              [torch.norm(p.grad.detach()) for p in model.parameters()]
          )
      )
      ret["loss"] = loss.item()
      ret["gnorm"] = gnorm.item()
  
      if ret["mem_0"] > 0:
          ret["mem_peak"] = torch.cuda.max_memory_allocated()
  
      return ret
  
imports: import pytest
  import torch
  import torch.nn
  from torch.utils.checkpoint import checkpoint
  from fairscale.fair_dev.testing.testing import skip_if_no_cuda
  from fairscale.internal import torch_version
  from fairscale.nn.checkpoint.checkpoint_activations import checkpoint_wrapper, disable_checkpointing
  from fairscale.nn.misc import FlattenParamsWrapper
  from fairscale.nn.misc import checkpoint_wrapper
label_desc: Use image node_id 2 for calling a global function with example usage: get_loss_and_gnorm(model, input) and returns: ret