HF pre-trained model for inference on multiple GPUs - parallel-processing

Question

I am trying to run inference from a pre-trained HF model on 4 GPUs. Do I need PyTorch Dataloader, or I can use any python function? I used pandas chunk split. My idea is to perform data parallelisation, split the data among 4 GPUs and send the pipeline HF_PIPELINE over each GPUs. Below is my my script. Could you suggest what I am doing wrong?
chunksize= 100
filename = 'dataset.csv'
def run_inference(world_size):
dist.init_process_group("gloo", rank=rank, world_size=world_size)
df_chunk_list = []
with pd.read_csv(filename, chunksize=chunksize) as reader:
for chunk in reader:
print(chunk.head(2))
chunk['prediction'] = list(map(HF_PIPELINE, chunk['text']))
df_chunk_list.append(chunk['prediction'])
print(df_chunk_list)
def main():
world_size = torch.cuda.device_count()
mp.spawn(run_inference,
args=(world_size),
nprocs=world_size,
join=True)
if __name__=="__main__":
main())
However I am getting this error
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/opt/conda/lib/python3.7/multiprocessing/spawn.py", line 105, in spawn_main
exitcode = _main(fd)
File "/opt/conda/lib/python3.7/multiprocessing/spawn.py", line 115, in _main
self = reduction.pickle.load(from_parent)
AttributeError: Can't get attribute 'run_inference' on <module '__main__' (built-in)>
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "/opt/conda/lib/python3.7/multiprocessing/spawn.py", line 105, in spawn_main
exitcode = _main(fd)
File "/opt/conda/lib/python3.7/multiprocessing/spawn.py", line 115, in _main
self = reduction.pickle.load(from_parent)
AttributeError: Can't get attribute 'run_inference' on <module '__main__' (built-in)>