Wrap the training loop you already have in one function and run it across sites. Iterate faster, ship it on time
Inside one training round
Each worker gets its own slice of the dataset. Nothing else about your data changes.
The body of the loop stays. It moves into one function that returns the weights and the sample count.
model = TinyCNN(num_classes=10) for _ in range(EPOCHS): train_one_epoch(model, DATA_DIR) torch.save(model.state_dict(), "out.pt")def fl_train(model, env, out_dir) -> dict: # your optimizer, your loss, your data for _ in range(env["EPOCHS"]): train_one_epoch(model, env["DATA_ROOT"]) return {"weights": model.state_dict(), "samples": n}Billed per GPU-minute.