michael/modrec-workflow
1
0
Fork 0
forked from qoherent/modrec-workflow
Code Pull Requests Workflows Activity

changed back to gpu

Browse Source
This commit is contained in:
Liyu Xiao 2025-05-26 14:38:18 -04:00
parent 92a0ed11e4
commit d518c633b7
1 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
data/training/train.py
View File

@ -119,11 +119,12 @@ def train_model():
)
)
checkpoint_callback = ModelCheckpoint(
dirpath=checkpoint_dir,
filename=checkpoint_filename,
save_top_k=1,
verbose=False, # Disable checkpoint verbose output
verbose=True, # Disable checkpoint verbose output
monitor="val_acc",
mode="max",
enable_version_counter=False,
@ -132,10 +133,10 @@ def train_model():
trainer = L.Trainer(
max_epochs=epochs,
callbacks=[checkpoint_callback],
accelerator="cpu",
accelerator="gpu",
devices=1,
benchmark=True,
precision="bf16-mixed",
precision="16-mixed",
logger=False,
enable_progress_bar=False, # Disable all progress bars
enable_model_summary=False, # Disable model summary