zc/mlp-project
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Disable img writer

Browse source
This commit is contained in:
Zhengyi Chen 2024-03-05 21:27:27 +00:00
parent 524ee03187
commit 733b3c5387
1 changed files with 3 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
train.py
View file

@ -335,10 +335,9 @@ def valid_one_epoch(test_loader, model, device, epoch, args):
mse += diff ** 2
if i % 5 == 0:
if isinstance(model, STNet_VisionTransformerGAP):
with torch.no_grad():
img_xformed = model.stnet(img).to("cpu")
xformed.append(img_xformed)
# with torch.no_grad():
# img_xformed = model.stnet(img).to("cpu")
# xformed.append(img_xformed)
print("[valid_one_epoch] {} | Gt {:.2f} Pred {:.4f} |".format(
fname[0],
torch.sum(gt_count_whole).item(),