Upload train.py with huggingface_hub

train.py

@@ -106,6 +106,11 @@ def train(args):
     print(f"[train] model params = {sum(p.numel() for p in YKDiff(cfg).parameters())/1e6:.1f}M")
 
     model = YKDiff(cfg).to(device)
+    # Train in bf16 weights so a large model fits a 22 GB L4: fp32 weights would
+    # need ~21.6 GB just for params+grads+Adam states (OOM). bf16 halves that.
+    if device.type in ("cuda", "xla"):
+        model = model.to(torch.bfloat16)
+        print("[train] using bfloat16 weights")
     n_params = sum(p.numel() for p in model.parameters())
     print(f"[train] allocated params = {n_params/1e6:.1f}M")