a-One-Fan/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

Add tag that can be applied to prevent parameter training

Browse Source
This commit is contained in:
James Betker 2020-10-06 20:39:49 -06:00
parent 2f2e3f33f8
commit 1e415b249b
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
codes/models/ExtensibleTrainer.py
View File

@ -192,7 +192,7 @@ class ExtensibleTrainer(BaseModel):
if net_enabled: if net_enabled:
enabled += 1 enabled += 1
for p in net.parameters(): for p in net.parameters():
if p.dtype != torch.int64 and p.dtype != torch.bool: if p.dtype != torch.int64 and p.dtype != torch.bool and not hasattr(p, "DO_NOT_TRAIN"):
p.requires_grad = net_enabled p.requires_grad = net_enabled
else: else:
p.requires_grad = False p.requires_grad = False