Aobangaming commited on
Commit
156abfb
·
verified ·
1 Parent(s): d9c0069

Update modeling_lightning.py

Browse files
Files changed (1) hide show
  1. modeling_lightning.py +2 -2
modeling_lightning.py CHANGED
@@ -2,7 +2,7 @@
2
  import torch
3
  import torch.nn as nn
4
 
5
- from transformers import PreTrainedModel, PretrainedConfig
6
  from transformers.modeling_outputs import CausalLMOutput
7
 
8
  import torch
@@ -190,7 +190,7 @@ class LightningConfig(PretrainedConfig):
190
  self.max_seq_len = max_seq_len
191
 
192
 
193
- class LightningForCausalLM(PreTrainedModel):
194
 
195
  config_class = LightningConfig
196
  base_model_prefix = "lightning"
 
2
  import torch
3
  import torch.nn as nn
4
 
5
+ from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
6
  from transformers.modeling_outputs import CausalLMOutput
7
 
8
  import torch
 
190
  self.max_seq_len = max_seq_len
191
 
192
 
193
+ class LightningForCausalLM(PreTrainedModel, GenerationMixin):
194
 
195
  config_class = LightningConfig
196
  base_model_prefix = "lightning"