config.json

{
  "architecture": "ConditionalMultiTraceUNet",
  "model": {
    "input_size": 32,
    "in_channels": 3,
    "num_classes": 10,
    "base_channels": 69,
    "channel_mult": [
      1,
      2,
      2,
      4
    ],
    "num_res_blocks": 1,
    "cond_dim": 312,
    "fourier_dim": 64,
    "attn_resolutions": [
      8
    ]
  },
  "objective": "multi_trace_imf",
  "multi_trace_imf": {
    "lambda_diag": 1.0,
    "max_nfe": 5,
    "jvp_precision": "bf16",
    "channels_last": true
  },
  "dataset": "cifar10",
  "weights_source": "ema",
  "max_nfe": 5,
  "training_steps": 50000,
  "trainable_parameters": 14992182,
  "pixel_range": [
    0,
    1
  ],
  "sampling_steps": 1
}