Laurent/REVA-QCAV
1
0
Fork 0
mirror of https://github.com/Laurent2916/REVA-QCAV.git synced 2024-11-09 15:02:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: typo

Browse source 
Former-commit-id: 0336f7e3a3809f6501dec4ca11a9ae1f3be44f29 [formerly 1bd5e48616d48d52e13067ff505a94162acf9e6f]
Former-commit-id: 471e9d11a26b14f6b878f9735d3e9c40eb10161b
This commit is contained in:
Laurent Fainsin 2022-06-30 14:43:41 +02:00
parent 24df16a612
commit 3056ebc498
1 changed files with 4 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
src/train.py
View file

@ -43,7 +43,7 @@ def main():
)
# create device
device = torch.device(wandb.config.device)
device = torch.device(wandb.config.DEVICE)
# enable cudnn benchmarking
torch.backends.cudnn.benchmark = wandb.config.BENCHMARK
@ -112,7 +112,7 @@ def main():
# save model.onxx
dummy_input = torch.randn(
1, wandb.config.n_channels, wandb.config.IMG_SIZE, wandb.config.IMG_SIZE, requires_grad=True
1, wandb.config.N_CHANNELS, wandb.config.IMG_SIZE, wandb.config.IMG_SIZE, requires_grad=True
).to(device)
torch.onnx.export(net, dummy_input, "model.onnx")
artifact = wandb.Artifact("onnx", type="model")
@ -132,8 +132,8 @@ def main():
# Training round
for step, (images, true_masks) in enumerate(train_loader):
assert images.shape[1] == net.N_CHANNELS, (
f"Network has been defined with {net.N_CHANNELS} input channels, "
assert images.shape[1] == net.n_channels, (
f"Network has been defined with {net.n_channels} input channels, "
f"but loaded images have {images.shape[1]} channels. Please check that "
"the images are loaded correctly."
)