Skip to content
Snippets Groups Projects
Commit d3ee0051 authored by Schneider Leo's avatar Schneider Leo
Browse files

fix : dataloader batched

parent f5fd0499
No related branches found
No related tags found
No related merge requests found
...@@ -19,6 +19,7 @@ def train_duo(model, data_train, optimizer, loss_function, epoch): ...@@ -19,6 +19,7 @@ def train_duo(model, data_train, optimizer, loss_function, epoch):
param.requires_grad = True param.requires_grad = True
for imaer,imana, img_ref, label in data_train: for imaer,imana, img_ref, label in data_train:
print(imaer.size())
label = label.long() label = label.long()
if torch.cuda.is_available(): if torch.cuda.is_available():
imaer = imaer.cuda() imaer = imaer.cuda()
...@@ -46,6 +47,7 @@ def test_duo(model, data_test, loss_function, epoch): ...@@ -46,6 +47,7 @@ def test_duo(model, data_test, loss_function, epoch):
param.requires_grad = False param.requires_grad = False
for imaer,imana, img_ref, label in data_test: for imaer,imana, img_ref, label in data_test:
print(imaer.size())
label = label.long() label = label.long()
if torch.cuda.is_available(): if torch.cuda.is_available():
imaer = imaer.cuda() imaer = imaer.cuda()
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment