#main.py
from train import train

#some preparatory code here
if __name__ == '__main__':
    train(x)

##########################
##########################
#train.py

def train(x):
    #initial data processing
    loader = DataLoader(dataset=_DATASET, batch_size=batch_size, shuffle=False, num_workers=4)
    #data post processing

This one should work. Would you please upload all your scripts so that I can have a look?