diag_hip.h 180 Bytes
Newer Older
quyuanhao123's avatar
quyuanhao123 committed
1
2
3
4
5
6
#pragma once

#include "../extensions.h"

torch::Tensor non_diag_mask_cuda(torch::Tensor row, torch::Tensor col,
                                 int64_t M, int64_t N, int64_t k);