小编Ana*_*das的帖子

我现在有更新的代码如下：

    # Hyperparameters
random_seed = 123
learning_rate = 0.01
num_epochs = 10
batch_size = 128

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

对于范围内的纪元（num_epochs）：模型= resnet34.train（）对于batch_idx，枚举（train_generator）中的（特征，目标）：

    features = features.to(device)
    targets = targets.to(device)
        
    ### FORWARD AND BACK PROP
    logits = model(features)
    cost = torch.nn.functional.cross_entropy(logits, targets)
    optimizer.zero_grad()
    
    cost.backward()
    
    ### UPDATE MODEL PARAMETERS
    optimizer.step()
    
    ### LOGGING
    if not batch_idx % 50:
        print ('Epoch: %03d/%03d | Batch %03d/%03d | Cost: %.4f' 
               %(epoch+1, num_epochs, batch_idx, 
                 len(datagen)//batch_size, cost))

model = model.eval() # eval mode to prevent upd. batchnorm …