code: def assert_correct_dimensions(sample, probs):
      ps_shape = list(probs.data.size())
      sample_shape = list(sample.shape)
      assert_equal(sample_shape, ps_shape)
  
file_path: pyro/tests/distributions/test_one_hot_categorical.py
imports: from unittest import TestCase
  import numpy
  import pytest
  import torch
  import pyro.distributions
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: assert_correct_dimensions(sample, probs) without return types