debug_hist.py

import torch
from problems.histogram.submission import custom_kernel, LENGTH, NUM_CHANNELS, NUM_BINS

def test():
    print("Testing custom_kernel...")
    length = LENGTH
    num_channels = NUM_CHANNELS
    num_bins = NUM_BINS
    
    print(f"Input: [{length}, {num_channels}], uint8")
    array = torch.randint(0, num_bins, (length, num_channels), dtype=torch.uint8).cuda()
    
    try:
        print("Calling custom_kernel...")
        out = custom_kernel((array, num_bins))
        print("Kernel finished.")
        print("Output shape:", out.shape)
        print("Output dtype:", out.dtype)
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    test()