submission 545146

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
import torch.nn.functional as F
from task import input_t, output_t

# Deterministic for correctness, but with benchmark to find fastest algorithm
torch.backends.cudnn.allow_tf32 = False
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True  # Find fastest deterministic algo

def custom_kernel(data: input_t) -> output_t:
    input_tensor, kernel, output = data
    result = F.conv2d(input_tensor, kernel, stride=1, padding=0)
    output[...] = result
    return output