Are GPU to Host copies also affected by Pinned Memory? I was wondering if we could collect all the results in GPU and transfer it one shot to CPU.
As you said, these copies didn’t affect my run time significantly. My network is taking 250ms. But I was wondering if there was any better way to do this.
Right now I have a torch Tensor pre allocated with the number of elements and in each iteration, I am indexing from n to n + batch_size on that tensor and storing the values. Hope I am not doing anything wrong by that.
def predict(self, dataloader):
    num_elements = len(dataloader.dataset)
    num_batches = len(dataloader)
    batch_size = dataloader.batch_size
    predictions = torch.zeros(num_elements)
    for i, batch in enumerate(dataloader):
        start = i*batch_size
        end = start + batch_size
        if i == num_batches - 1:
            end = num_elements
        pred, output = self.step(batch)
        predictions[start:end] = pred