Run dir : output/_smoke_test_1gpu
Log file: output/_smoke_test_1gpu/train.log
GPU: NVIDIA GeForce RTX 5090 | VRAM: 31.4 GiB | PyTorch: 2.11.0+cu130

Final Configuration:
  Paths:
    transformer_path             weights/flux2_dev_fp8mixed.safetensors
    vae_path                     weights/flux2-vae.safetensors
    controlnet_path              weights/FLUX.2-dev-Fun-Controlnet-Union-2602.safetensors
    dataset_dir                  dataset
    color_map_path               configs/color_map.json
    output_dir                   output/_smoke_test_1gpu
    text_encoder_path            weights/mistral_3_small_flux2_fp8.safetensors
    precomputed_embeddings       output/text_embeddings_global.pt
  Model:
    image_size                   1024
    num_classes                  6
    control_in_dim               3072
    fusion_dim                   768
    num_fusion_blocks            3
    num_heads                    12
    num_fourier_bands            32
    boundary_threshold           0.1
  Training:
    num_epochs                   1
    batch_size                   4
    learning_rate                0.0003
    weight_decay                 0.01
    max_grad_norm                1.0
    grad_accum_steps             4
    guidance_scale               3.5
    num_workers                  0
  Text Encoder:
    text_seq_len                 512
    text_dim                     15360
  Logging:
    log_interval                 1
    save_every_n_epochs          5
    val_every_n_epochs           1
  WandB:
    wandb_entity                 
    wandb_project                _smoke_test_1gpu
  Resume:
    resume_from                  (not set)
[MEM @ pre-flight] RAM: 11.8/188.5 GiB (6.3%)  |  VRAM: 0.0/31.4 GiB (0.0%)

============================================================
  [1/8] Text Embeddings
============================================================
  Composed prompt (575 chars):
    Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite imagery, natural daylight, buildings: detaile...

=== Precomputing Global Text Embedding ===
  Prompt: Aerial top-down satellite view of American urban area, Google Earth style, 8k resolution, photorealistic satellite image...
  Loading safetensors: weights/mistral_3_small_flux2_fp8.safetensors
  Building tokenizer from embedded tekken_model ...
  Detected: 30 layers, hidden=5120, heads=32, kv_heads=8, ffn=32768, vocab=131072
  Initialising MistralModel (30 layers)...
  Dequantizing FP8 weights → bf16 ...
Traceback (most recent call last):
  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 1297, in <module>
    main()
  File "/home/xg_wang_group/SynthUrbanSAT/train_script.py", line 731, in main
    precompute_single_prompt_embeddings(
  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 449, in precompute_single_prompt_embeddings
    text_encoder, tokenizer = load_text_encoder(model_path, device, dtype)
                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 115, in load_text_encoder
    return _load_from_safetensors(model_path, device, dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xg_wang_group/SynthUrbanSAT/scripts/text_encoder.py", line 213, in _load_from_safetensors
    text_encoder = text_encoder.to(device=device, dtype=dtype).eval()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/transformers/modeling_utils.py", line 3574, in to
    return super().to(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1384, in to
    return self._apply(convert)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
    module._apply(fn)
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
    module._apply(fn)
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 934, in _apply
    module._apply(fn)
  [Previous line repeated 1 more time]
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 965, in _apply
    param_applied = fn(param)
                    ^^^^^^^^^
  File "/home/xg_wang_group/miniconda/envs/flux_train/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1370, in convert
    return t.to(
           ^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 320.00 MiB. GPU 0 has a total capacity of 31.36 GiB of which 208.56 MiB is free. Including non-PyTorch memory, this process has 31.13 GiB memory in use. Of the allocated memory 30.65 GiB is allocated by PyTorch, and 1.45 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://docs.pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf)