File tree Expand file tree Collapse file tree 1 file changed +3
-6
lines changed Expand file tree Collapse file tree 1 file changed +3
-6
lines changed Original file line number Diff line number Diff line change 55
55
with open (args .checkpoint , 'rb' ) as f :
56
56
safe_globals = [
57
57
PositionalEncoding ,
58
+ RNNModel ,
58
59
TransformerModel ,
59
60
torch .nn .functional .relu ,
60
61
torch .nn .modules .activation .MultiheadAttention ,
64
65
torch .nn .modules .linear .NonDynamicallyQuantizableLinear ,
65
66
torch .nn .modules .normalization .LayerNorm ,
66
67
torch .nn .modules .sparse .Embedding ,
67
- torch .nn .modules .transformer .TransformerEncoder ,
68
- torch .nn .modules .transformer .TransformerEncoderLayer ,
69
- RNNModel ,
70
- torch .nn .modules .dropout .Dropout ,
71
- torch .nn .modules .linear .Linear ,
72
68
torch .nn .modules .rnn .GRU ,
73
69
torch .nn .modules .rnn .LSTM ,
74
70
torch .nn .modules .rnn .RNN ,
75
- torch .nn .modules .sparse .Embedding ,
71
+ torch .nn .modules .transformer .TransformerEncoder ,
72
+ torch .nn .modules .transformer .TransformerEncoderLayer ,
76
73
]
77
74
with torch .serialization .safe_globals (safe_globals ):
78
75
model = torch .load (f , map_location = device )
You can’t perform that action at this time.
0 commit comments