I am new to PyTorch and want to efficiently evaluate among others F1 during my Training and my Validation Loop.

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.

```
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 > 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')
```

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?

I also posted this question on stackexchange, but thought that this forum is probably even better

As I am new to PyTorch, I am very grateful for any hints and feedback