How to use custom collate_fn by update defalut_collate_fn_map

torch.utils.data — PyTorch 2.3 documentation

In the page, it says default_collate_fn_map.update({torch.Tensor: collate_tensor_fn}) can be used to update defalut_collate. For my case, I want to add default_collate_fn_map.update({PIL.Image: lambda x: return x}). However, I cannot change default_collate_fn_map because it’s defined in torch/utils/data/_utils/collate.py. Any help?

import torch
torch.utils.data._utils.collate.default_collate_fn_map.update({})

Seems the default_collate_fn_map is located in a internal package _utils.
It’s not exposed to user as the default_collate api.