"src/vscode:/vscode.git/clone" did not exist on "a4bc8454783dbae3be9cf840f074b961a558aba5"
gpu_decoder.h 435 Bytes
Newer Older
Prabhat Roy's avatar
Prabhat Roy committed
1
2
3
4
5
6
7
8
9
10
#include <torch/custom_class.h>
#include <torch/torch.h>
#include "decoder.h"
#include "demuxer.h"

class GPUDecoder : public torch::CustomClassHolder {
 public:
  GPUDecoder(std::string, int64_t);
  ~GPUDecoder();
  torch::Tensor decode();
11
  c10::Dict<std::string, c10::Dict<std::string, double>> get_metadata() const;
Prabhat Roy's avatar
Prabhat Roy committed
12
13
14
15
16
17
18
19

 private:
  Demuxer demuxer;
  CUcontext ctx;
  Decoder decoder;
  int64_t device;
  bool initialised = false;
};