Customized collate_fn gives False pin_memory in the official example

AlexisW · May 11, 2019, 7:47am

pytorch/pytorch/blob/master/torch/utils/data/dataloader.py#L112


The default memory pinning logic only recognizes Tensors and maps and iterables
containg Tensors.  By default, if the pinning logic sees a batch that is a custom type
(which will occur if you have a ``collate_fn`` that returns a custom batch type),
or if each element of your batch is a custom type, the pinning logic will not
recognize them, and it will return that batch (or those elements)
without pinning the memory.  To enable memory pinning for custom batch or data types,
define a ``pin_memory`` method on your custom type(s).


Example::


    class SimpleCustomBatch:
        def __init__(self, data):
            transposed_data = list(zip(*data))
            self.inp = torch.stack(transposed_data[0], 0)
            self.tgt = torch.stack(transposed_data[1], 0)


        def pin_memory(self):
            self.inp = self.inp.pin_memory()
            self.tgt = self.tgt.pin_memory()
            return self

Running this piece of code gives me

False
False
False
False
False
False
False
False
False
False

Is that normal?

SimonW · May 11, 2019, 10:14pm

it’s normal if you are using old pytorch or does not have cuda