Move the model to gpu
parent
f1e06e97ad
commit
acd5250c12
2
ft.py
2
ft.py
|
@ -21,6 +21,8 @@ class FT:
|
||||||
self.model_name
|
self.model_name
|
||||||
)
|
)
|
||||||
|
|
||||||
|
self.model.to(self.device)
|
||||||
|
|
||||||
# set up optimizer
|
# set up optimizer
|
||||||
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)
|
self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=1e-5)
|
||||||
self.scheduler = ReduceLROnPlateau(
|
self.scheduler = ReduceLROnPlateau(
|
||||||
|
|
Loading…
Reference in New Issue