kernel
File size: 336 Bytes
2595c46
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
#pragma once

#include <torch/all.h>

namespace megablocks {

// Public interface function for constructing indices from padded bins
void indices(torch::Tensor padded_bins,
             int block_size,
             int output_block_rows,
             int output_block_columns,
             torch::Tensor out);

} // namespace megablocks