ZinebSN commited on
Commit
206fa83
1 Parent(s): a535e2c

Update t5_model.py

Browse files
Files changed (1) hide show
  1. t5_model.py +1 -0
t5_model.py CHANGED
@@ -1,5 +1,6 @@
1
  import pytorch_lightning as pl
2
  from transformers import AutoModelForSeq2SeqLM
 
3
 
4
  class T5(pl.LightningModule):
5
  def __init__(self, lr=5e-5, num_train_epochs=15, warmup_steps=1000):
 
1
  import pytorch_lightning as pl
2
  from transformers import AutoModelForSeq2SeqLM
3
+ import torch
4
 
5
  class T5(pl.LightningModule):
6
  def __init__(self, lr=5e-5, num_train_epochs=15, warmup_steps=1000):