module.h 376 Bytes
Newer Older
dengjb's avatar
update  
dengjb committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#pragma once

#include <map>
#include "struct.h"
#include "NvInfer.h"
using namespace nvinfer1;

namespace fastrt {

    class Module {
    public:
        Module() = default;
        virtual ~Module() = default;

        virtual ILayer* topology(INetworkDefinition *network, 
            std::map<std::string, Weights>& weightMap, 
            ITensor& input) = 0; 
    };

}