You mean including a method like:
def channels_last(input):
    return(input.contiguous(memory_format=torch.channels_last))

and then putting that inside transforms.Lambda? That gives the same “can’t pickle” error AttributeError: Can't pickle local object 'main.<locals>.channels_last' if it’s included locally, and if it’s located globally then it appears to get stuck on something when inside the training loop–no error is thrown but it doesn’t complete an epoch (and system utilisation drops to nothing).
Writing a custom dataset would entail rewriting the whole functionality of ImageFolder, which I was hoping to avoid.