Function main

samples/python/onnx_custom_plugin/test_custom_hardmax_plugin.py:71–82

()

    return output

def main():
    load_plugin_lib()
    for num_dims in range(1, 8):
        for axis in range(-num_dims, num_dims):
            shape = np.random.randint(1, 4, size=num_dims)
            arr = np.random.rand(*shape)
            arr = (arr - 0.5) * 200
            engine = make_trt_network_and_engine(shape, axis)
            res1 = hardmax_reference_impl(arr, axis)
            res2 = custom_plugin_impl(arr, engine).reshape(res1.shape)
            assert np.all(res1 == res2), f"Test failed for shape={shape}, axis={axis}"
    print("Passed")

if __name__ == '__main__':
    main()