code: def mock_safe_multinomial(
      input: torch.Tensor, num_samples: int
  ) -> torch.Tensor:
      """Return non deterministic indexes to mock safe_multinomial
  
      Args:
          input: tensor of shape [B, n] containing non-negative values;
                  rows are interpreted as unnormalized event probabilities
                  in categorical distributions.
          num_samples: number of samples to take.
  
      Returns:
          Tensor of shape [B, num_samples]
      """
      batch_size = input.shape[0]
      return torch.arange(num_samples).repeat(batch_size, 1).to(DEVICE)
  
file_path: pytorch3d/tests/implicitron/models/test_overfit_model.py
imports: import unittest
  from typing import Any, Dict
  from unittest.mock import patch
  import torch
  from pytorch3d.implicitron.models.generic_model import GenericModel
  from pytorch3d.implicitron.models.overfit_model import OverfitModel
  from pytorch3d.implicitron.models.renderer.base import EvaluationMode
  from pytorch3d.implicitron.tools.config import expand_args_fields
  from pytorch3d.renderer.cameras import look_at_view_transform, PerspectiveCameras
label_desc: Use image node_id 2 for calling a global function with example usage: mock_safe_multinomial(input, num_samples) and returns: torch