Machine-Learning
/
aiuNN
1
0
Fork 0
Code Issues 2 Pull Requests Actions Packages Projects Releases Wiki Activity

develop #4

Merged
Fabel merged 103 commits from develop into main 2025-03-01 21:47:17 +00:00
Conversation 0 Commits 103 Files Changed 13 +1
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit e69d0e90ec - Show all commits

1
src/aiunn/finetune.py
View File

@ -135,6 +135,7 @@ def finetune_model(model: nn.Module, datasets: list[str], batch_size=1, epochs=1
with autocast(device_type="cuda"): with autocast(device_type="cuda"):
if use_checkpoint: if use_checkpoint:
# Use checkpointing if requested. # Use checkpointing if requested.
low_res = batch['low_res'].to(device).requires_grad_()
features = checkpoint(lambda x: model(x), low_res) features = checkpoint(lambda x: model(x), low_res)
else: else:
features = model(low_res) features = model(low_res)