title: Efficient metrics evaluation in PyTorch
body: <p>I am new to PyTorch and want to efficiently evaluate among others F1 during my Training and my Validation Loop.</p>
  
  <p>So far, my approach was to calculate the predictions on GPU, then push them to CPU and append them to a vector for both Training and Validation. After Training and Validation, I would evaluate both for each epoch using sklearn. However, profiling my code it showed, that pushing to cpu is quite a bottleneck. </p>
  
  <pre class="lang-py prettyprint-override"><code>for epoch in range(n_epochs):
      model.train()
      avg_loss = 0
      avg_val_loss = 0
      train_pred = np.array([])
      val_pred = np.array([])
      # Training loop (transpose X_batch to fit pretrained (features, samples) style)
      for X_batch, y_batch in train_loader:
          scores = model(X_batch)
          y_pred = F.softmax(scores, dim=1)
          train_pred = np.append(train_pred, self.get_vector(y_pred.detach().cpu().numpy()))
  
          loss = loss_fn(scores, self.get_vector(y_batch))
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
          avg_loss += loss.item() / len(train_loader)
  
      model.eval()
      # Validation loop
      for X_batch, y_batch in val_loader:
          with torch.no_grad():
              scores = model(X_batch)
              y_pred = F.softmax(scores, dim=1)
              val_pred = np.append(val_pred, self.get_vector(y_pred.detach().cpu().numpy()))
              loss = loss_fn(scores, self.get_vector(y_batch))
              avg_val_loss += loss.item() / len(val_loader)
  
      # Model Checkpoint for best validation f1
      val_f1 = self.calculate_metrics(train_targets[val_index], val_pred, f1_only=True)
      if val_f1 &gt; best_val_f1:
          prev_best_val_f1 = best_val_f1
          best_val_f1 = val_f1
          torch.save(model.state_dict(), self.PATHS['xlm'])
          evaluated_epoch = epoch
  
      # Calc the metrics
      self.save_metrics(train_targets[train_index], train_pred, avg_loss, 'train')
      self.save_metrics(train_targets[val_index], val_pred, avg_val_loss, 'val')
  </code></pre>
  
  <p>I am certain there is a more efficient way to 
  a) store the predictions without having to push them to cpu each batch. b) calculate the metrics on GPU directly?</p>
  
  <p>As I am new to PyTorch, I am very grateful for any hints and feedback :)</p>