| 4 | #include<BaseModule.h> |
| 5 | |
| 6 | class MLP: public torch::nn::Module{ |
| 7 | public: |
| 8 | MLP(int in_features, int out_features); |
| 9 | torch::Tensor forward(torch::Tensor x); |
| 10 | private: |
| 11 | int mid_features[3] = {32,64,128}; |
| 12 | LinearBnRelu ln1{nullptr}; |
| 13 | LinearBnRelu ln2{nullptr}; |
| 14 | LinearBnRelu ln3{nullptr}; |
| 15 | torch::nn::Linear out_ln{nullptr}; |
| 16 | }; |