Collect Threads after Trainer.fit

[42elenz]

Hello!

I want to collect multipe threads after my trainer.fit because I am running multiple models one after the other in my script. But PL spawns multiple threads not just in the trainer-function but for the whole script (what I acutally find a bit weird though and didn't knew before).

Is there a way to "collect" threads after the trainer.fit?
Also I am saving best models and this is also broken due to the mutlipe threads (suddenly there are two folders).

Can somebody help me with that? Would be awesome :)

def pretrain(args, name, logdir, train_loader, val_loader, train_dataset, val_dataset, date_tag, fold=''):
    set_corruption_for_datasets(args.mri_corruption_pretrain, args.quest_corruption_pretrain, args.val_uncorrupted_pretrain, train_dataset, val_dataset)
    if args.wandb:
        name = remove_fold_suffix(name)
        name = f"{name}_{date_tag}"
        wandb_logger_pt = WandbLogger(project=args.project_name, name=name, reinit=True, tags=fold)
    monitor_metric = args.pretraining_metric
    maximize_minimize = args.pre_maximize_minimize
    checkpoint_callback, csv_logger_callback, args = set_log_dirs(logdir, monitor_metric,maximize_minimize, args, wandb_logger=wandb_logger_pt)
    model = choose_model(args)
    if wandb_logger_pt:
        wandb_logger_pt.watch(model, log='all')
    trainer = Trainer(
        max_epochs=args.training.epochs,
        logger=wandb_logger_pt,
        #log_every_n_steps = args.logging.log_every_n_steps,
        callbacks=[checkpoint_callback],
        accelerator="gpu",
        precision='16-mixed',
        strategy="ddp_find_unused_parameters_true"
    )
    trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader,)
    trainer.teardown()
    print(f"Finished training, Threading: {threading.enumerate()}")
    clean_gpu_memory()
    found_validation_ids_to_validation_df(val_dataset, args, logdir)