@triton.jit
def _ld_mse_fwd_kernel(loss_ptr, input_ptr, target_ptr, loss_row_stride,
    input_row_stride, target_row_stride, n_rows, n_cols, BLOCK_SIZE: tl.
    constexpr):
    pid = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols
    input_ptrs = input_ptr + pid * input_row_stride + col_offsets
    target_ptrs = target_ptr + pid * target_row_stride + col_offsets
    input = tl.load(input_ptrs, mask=mask, other=0.0)
    target = tl.load(target_ptrs, mask=mask, other=0.0)
    loss = tl.sum((input - target) * (input - target)) / n_cols
    loss_ptrs = loss_ptr + pid
    tl.store(loss_ptrs, loss, mask=pid < n_rows)
