File size: 662 Bytes
4c954ae
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
// #pragma once
#include <torch/extension.h>

std::tuple<at::Tensor, at::Tensor, at::Tensor> lsencode_cuda(
    const at::Tensor& lines,
    const int input_height,
    const int input_width,
    const int height,
    const int width,
    const int num_lines);
    
std::tuple<at::Tensor,at::Tensor,at::Tensor> encodels(
    const at::Tensor& lines,
    const int input_height,
    const int input_width,
    const int height,
    const int width,
    const int num_lines)
{
    return lsencode_cuda(lines,
                    input_height,
                    input_width,
                    height,
                    width,
                    num_lines);
}