Listing: Listing

def replay_late_branch(
    early_state, early_masks, loaders,
    late_fraction, training_seed,
    epochs, learning_rate
):
    # CPU implementation matching the supplied DEVICE.
    with torch.random.fork_rng(devices=[]):
        torch.manual_seed(int(training_seed))
        model = TaskGatedNetwork().to("cpu")
        model.load_state_dict(early_state)

        manager = TaskGatedPruningManager(model)
        manager.masks = {
            name: mask.clone().to("cpu")
            for name, mask in early_masks.items()
        }
        manager.apply_masks()
        manager.prune_fraction_of_remaining(late_fraction)

        local_loaders = dict(loaders)
        for offset, task in enumerate(("visual", "voice"), start=1):
            source = loaders[f"{task}_train"]
            generator = torch.Generator(device="cpu")
            generator.manual_seed(int(training_seed) + offset)
            local_loaders[f"{task}_train"] = DataLoader(
                source.dataset,
                batch_size=source.batch_size,
                shuffle=True,
                num_workers=0,
                generator=generator,
            )

        # Construction consumed global random numbers;
        # reset the local training stream explicitly.
        torch.manual_seed(int(training_seed))
        train_taskgated_with_curves(
            model,
            local_loaders,
            epochs=epochs,
            lr=learning_rate,
            pruning_manager=manager,
            verbose=False,
            track_test_accuracy=False,
        )

    return model, manager