| 41 | } |
| 42 | |
| 43 | static std::unique_ptr<layers::BertPooler> LoadPooler(NPZMapView npz, |
| 44 | DLDeviceType dev) { |
| 45 | NPZLoader params(std::move(npz), dev); |
| 46 | |
| 47 | return std::unique_ptr<layers::BertPooler>( |
| 48 | new layers::BertPooler(params["dense.weight"], params["dense.bias"])); |
| 49 | } |
| 50 | |
| 51 | struct BERTLayer { |
| 52 | explicit BERTLayer(NPZLoader params, int64_t n_heads) { |