Commit dceace67 authored by Martin Lebourdais's avatar Martin Lebourdais
Browse files

removal of the weights

parent 4533586e
......@@ -477,7 +477,7 @@ def train_epoch(model, epoch, training_loader, optimizer, log_interval, device):
"""
model.to(device)
model.train()
criterion = torch.nn.CrossEntropyLoss(reduction='mean',weight=torch.FloatTensor([0.1,0.9]).to(device))
criterion = torch.nn.CrossEntropyLoss(reduction='mean')
recall = 0.0
accuracy = 0.0
for batch_idx, (data, target) in enumerate(training_loader):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment