code: def test_categorical():
      probs = torch.as_tensor([0.2, 0.6, 0.2])
      dist = torch.distributions.Categorical(probs=probs)
      mode = mode_of_distribution(dist)
      assert mode.tolist() == 1
  
      probs = torch.as_tensor([0.6, 0.2, 0.2])
      dist = torch.distributions.Categorical(probs=probs)
      mode = mode_of_distribution(dist)
      assert mode.tolist() == 0
  
      probs = torch.as_tensor([[0.6, 0.2, 0.2], [0.2, 0.2, 0.6]])
      dist = torch.distributions.Categorical(probs)
      mode = mode_of_distribution(dist)
      assert mode.tolist() == [0, 2]
  
file_path: pfrl/tests/utils_tests/test_mode_of_distribution.py
imports: import math
  import torch
  from pfrl.testing import torch_assert_allclose
  from pfrl.utils.mode_of_distribution import mode_of_distribution
label_desc: Use image node_id 2 for calling a global function with example usage: test_categorical() without return types