Training with larger input sizes (e.g., 1280x1280) leads to NaN loss

[Ibrah-N]

Hi Sebastian,

First of all, thank you for releasing this great work — I’ve been experimenting with LINEA and it has been performing really well on my dataset.

So I have trained the model for the images of patches 1024x1024 without changing the any original repo configuration and it was giving a very good results, but I want to train the model for larger sizes like 2048 or may higher, to test if the model better generalize on larger images so I train for larger patches to gain more context level special features.

So to try the model for larger input size that starts giving me issues while updating all those below params with the rule of thumb (multiplying by 2). Below is details explanation of the configuration changes I made.

---

configs/linea/include/dataset.py
===== original configuration =====

data_aug_scales = [(640, 640)]
data_aug_max_size = 1333
data_aug_scales2_resize = [400, 500, 600]
data_aug_scales2_crop = [384, 600]

===== changed configuration =====

data_aug_scales = [(1280, 1280)]
data_aug_max_size = 2666
data_aug_scales2_resize = [800, 1000, 1200]
data_aug_scales2_crop = [768, 1200]

---

configs/linea/include/linea.py
===== original configuration =====

modelname = 'LINEA'
eval_spatial_size = (640, 640)
eval_idx = 5 # 6 decoder layers
num_classes = 2
...
...
...

===== changed configuration =====

modelname = 'LINEA'
eval_spatial_size = (1280, 1280)
eval_idx = 5 # 6 decoder layers
num_classes = 2
...
...
...

---

Model Training:

%cd /content/LINEA
!torchrun --nproc_per_node=1 main.py \
  --config_file configs/linea/linea_hgnetv2_l.py \
  --coco_path data/wireframe_processed \
  --resume weights/linea_hgnetv2_l.pth \
  --num_workers 2

---

OUTPUT LOGS:
/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py:4807: UserWarning: No device id is provided via init_process_group or barrier . Using the current device set by the user.
warnings.warn( warn only once
[rank0]:[W906 13:14:41.801293515 ProcessGroupNCCL.cpp:5023] [PG ID 0 PG GUID 0 Rank 0] using GPU 0 as device used by this process is currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can specify device_id in init_process_group() to force use of a particular device.
Initialized distributed mode...
Namespace(config_file='configs/linea/linea_hgnetv2_l.py', options=None, coco_path='data/wireframe_processed', device='cuda', seed=42, resume='', start_epoch=0, eval=False, num_workers=2, find_unused_params=False, world_size=1, rank=0, local_rank=None, amp=False, gpu=0, distributed=True, data_aug_scales=[(1280, 1280)], data_aug_max_size=2666, data_aug_scales2_resize=[800, 1000, 1200], data_aug_scales2_crop=[768, 1200], data_aug_scale_overlap=None, batch_size_train=2, batch_size_val=2, lr=0.00025, weight_decay=0.000125, betas=[0.9, 0.999], epochs=16, lr_drop_list=[11], clip_max_norm=0.1, save_checkpoint_interval=1, modelname='LINEA', eval_spatial_size=(1024, 1024), eval_idx=5, num_classes=2, pretrained=True, use_checkpoint=False, return_interm_indices=[1, 2, 3], freeze_norm=True, freeze_stem_only=True, hybrid_encoder='hybrid_encoder_asymmetric_conv', in_channels_encoder=[512, 1024, 2048], pe_temperatureH=20, pe_temperatureW=20, transformer_activation='relu', batch_norm_type='FrozenBatchNorm2d', masks=False, aux_loss=True, num_queries=1100, query_dim=4, num_feature_levels=3, dec_n_points=[4, 1, 1], dropout=0.0, pre_norm=False, use_dn=True, dn_number=300, dn_line_noise_scale=1.0, dn_label_noise_ratio=0.5, embed_init_tgt=True, dn_labelbook_size=2, match_unstable_error=True, set_cost_class=2.0, set_cost_lines=5.0, criterionname='LINEACRITERION', criterion_type='default', weight_dict={'loss_logits': 4, 'loss_line': 5}, losses=['labels', 'lines'], focal_alpha=0.1, matcher_type='HungarianMatcher', nms_iou_threshold=-1, use_ema=False, ema_decay=0.9997, ema_epoch=0, output_dir='output/linea_hgnetv2_l', backbone='HGNetv2_B4', param_dict_type='hgnetv2_b4', use_lab=False, feat_strides=[8, 16, 32], hidden_dim=256, dim_feedforward=1024, nheads=8, use_lmap=False, expansion=0.5, depth_mult=1.0, feat_channels_decoder=[256, 256, 256], dec_layers=6, num_select=300, reg_max=16, reg_scale=4, use_warmup=False, model_parameters=[{'params': '^(?=.backbone)(?!.norm|bn).$', 'lr': 1.25e-05}, {'params': '^(?=.(?:encoder|decoder))(?=.(?:norm|bn)).$', 'weight_decay': 0.0}])
Loading stage1
Loaded stage1 B4 HGNetV2 from local file.
building train_dataloader with batch_size=2...
loading annotations into memory...
Done (t=0.03s)
creating index...
index created!
building val_dataloader with batch_size=2...
loading annotations into memory...
Done (t=0.00s)
creating index...
index created!
Multi-scaling uses the following size: [768, 800, 832, 864, 896, 928, 960, 992, 1024, 1024, 1024, 1056, 1088, 1120, 1152, 1184, 1216, 1248, 1280]

------------------------------------- Calculate Flops Results -------------------------------------
Notations:
number of parameters (Params), number of multiply-accumulate operations(MACs),
number of floating-point operations (FLOPs), floating-point operations per second (FLOPS),
fwd FLOPs (model forward propagation FLOPs), bwd FLOPs (model backward propagation FLOPs),
default model backpropagation takes 2.00 times as much computation as forward propagation.

Total Training Params: 25.04 M
fwd MACs: 91.9696 GMACs
fwd FLOPs: 184.327 GFLOPS
fwd+bwd MACs: 275.909 GMACs
fwd+bwd FLOPs: 552.98 GFLOPS

{'flops': '184.327 GFLOPS', 'macs': '91.9696 GMACs', 'params': 25036882}
----------------------------------------- Start training ------------------------------------------
Epoch: [0] [ 0/422] eta: 0:27:19 lr: 0.000013 loss: 79.0159 (79.0159) loss_line: 3.3012 (3.3012) loss_line_0: 3.3077 (3.3077) loss_line_1: 3.3077 (3.3077) loss_line_2: 3.3012 (3.3012) loss_line_3: 3.3012 (3.3012) loss_line_4: 3.3012 (3.3012) loss_line_dn_0: 5.7845 (5.7845) loss_line_dn_1: 5.7915 (5.7915) loss_line_dn_2: 5.8021 (5.8021) loss_line_dn_3: 5.8139 (5.8139) loss_line_dn_4: 5.8252 (5.8252) loss_line_dn_5: 5.8357 (5.8357) loss_line_interm: 3.3012 (3.3012) loss_logits: 1.6663 (1.6663) loss_logits_0: 1.4373 (1.4373) loss_logits_1: 1.5000 (1.5000) loss_logits_2: 1.5760 (1.5760) loss_logits_3: 1.6247 (1.6247) loss_logits_4: 1.6522 (1.6522) loss_logits_dn_0: 1.6217 (1.6217) loss_logits_dn_1: 1.6513 (1.6513) loss_logits_dn_2: 1.7024 (1.7024) loss_logits_dn_3: 1.7407 (1.7407) loss_logits_dn_4: 1.7579 (1.7579) loss_logits_dn_5: 1.7574 (1.7574) loss_logits_interm: 1.3535 (1.3535) time: 3.8844 data: 1.3720 max mem: 6224
Loss is nan, stopping training
{'loss_line': tensor(0., device='cuda:0', grad_fn=), 'loss_line_0': tensor(0., device='cuda:0', grad_fn=), 'loss_line_1': tensor(0., device='cuda:0', grad_fn=), 'loss_line_2': tensor(0., device='cuda:0', grad_fn=), 'loss_line_3': tensor(0., device='cuda:0', grad_fn=), 'loss_line_4': tensor(0., device='cuda:0', grad_fn=), 'loss_line_dn_0': tensor(0., device='cuda:0', grad_fn=), 'loss_line_dn_1': tensor(0., device='cuda:0', grad_fn=), 'loss_line_dn_2': tensor(0., device='cuda:0', grad_fn=), 'loss_line_dn_3': tensor(0., device='cuda:0', grad_fn=), 'loss_line_dn_4': tensor(0., device='cuda:0', grad_fn=), 'loss_line_dn_5': tensor(0., device='cuda:0', grad_fn=), 'loss_line_interm': tensor(0., device='cuda:0', grad_fn=), 'loss_logits': tensor(0.6740, device='cuda:0', grad_fn=), 'loss_logits_0': tensor(1.3308, device='cuda:0', grad_fn=), 'loss_logits_1': tensor(0.4909, device='cuda:0', grad_fn=), 'loss_logits_2': tensor(0.5068, device='cuda:0', grad_fn=), 'loss_logits_3': tensor(0.4487, device='cuda:0', grad_fn=), 'loss_logits_4': tensor(0.5344, device='cuda:0', grad_fn=), 'loss_logits_dn_0': tensor(nan, device='cuda:0', grad_fn=), 'loss_logits_dn_1': tensor(nan, device='cuda:0', grad_fn=), 'loss_logits_dn_2': tensor(nan, device='cuda:0', grad_fn=), 'loss_logits_dn_3': tensor(nan, device='cuda:0', grad_fn=), 'loss_logits_dn_4': tensor(nan, device='cuda:0', grad_fn=), 'loss_logits_dn_5': tensor(nan, device='cuda:0', grad_fn=), 'loss_logits_interm': tensor(0.4288, device='cuda:0', grad_fn=)}
[rank0]:[W906 13:15:48.271268219 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
E0906 13:15:50.690000 21783 torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 0 (pid: 21790) of binary: /usr/bin/python3
Traceback (most recent call last):
File "/usr/local/bin/torchrun", line 10, in
sys.exit(main())
^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 357, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 901, in main
run(args)
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 892, in run
elastic_launch(
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 143, in call
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 277, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

main.py FAILED

Failures:
<NO_OTHER_FAILURES>

Root Cause (first observed failure):
[0]:
time : 2025-09-06_13:15:50
host : d51c69c24c24
rank : 0 (local_rank: 0)
exitcode : 1 (pid: 21790)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================"

[SebastianJanampa]

Hi @Ibrah-N,
Thanks for the kind words. About your problem, you should modify this line of code

LINEA/datasets/coco.py

Lines 123 to 147 in 475c5ce

	# update args from config files
	scales = args.data_aug_scales
	max_size = args.data_aug_max_size
	scales2_resize = args.data_aug_scales2_resize
	scales2_crop = args.data_aug_scales2_crop
	test_size = args.eval_spatial_size
	if image_set == 'train':
	return T.Compose([
	T.RandomSelect(
	T.RandomHorizontalFlip(),
	T.RandomVerticalFlip(),
	),
	T.RandomSelect(
	T.RandomResize(scales, max_size=max_size),
	T.Compose([
	T.RandomResize(scales2_resize),
	T.RandomSizeCrop(384, 600),
	T.RandomResize(scales, max_size=max_size),
	])
	),
	T.ColorJitter(),
	normalize,
	])

line 140 should be

                 T.RandomSizeCrop(*scales2_crop)

I would also recommend that you run datasets/coco.py to visualize how the augmentation techniques affect your dataset. You should only modify these lines

LINEA/datasets/coco.py

Lines 193 to 203 in 475c5ce

	dataset_debug = CocoDetection(
	'../data/wireframe_processed/val2017',
	'../data/wireframe_processed/annotations/lines_val2017.json',
	transforms=T.Compose([
	# T.RandomResize([400, 500, 600]),
	# T.RandomSizeCrop(384, 600),
	T.RandomResize([(640, 640)], max_size=1333),
	T.ToTensor()
	]),
	include_lmap=False
	)

You should add the augmentation techniques you want to use in your training, and add the path to your dataset.

[Ibrah-N]

Thank you @SebastianJanampa

Your help and quick reply is appreciated....

Previously it was raising error at the first epoch. But after making the changes you suggest -- the model starts training and the results of the model was also good but at the 10th epoch it again raised the same issues; logs is pasted. Your help and reply will greatly be considered.

---

Model Training:

%cd /content/LINEA
!torchrun --nproc_per_node=1 main.py \
  --config_file configs/linea/linea_hgnetv2_l.py \
  --coco_path data/wireframe_processed \
  --num_workers 2

---

Model Training Logs:

/content/LINEA`
/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py:4807: UserWarning: No device id is provided via `init_process_group` or `barrier `. Using the current device set by the user. 
  warnings.warn(  # warn only once
[rank0]:[W907 05:30:19.731165827 ProcessGroupNCCL.cpp:5023] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 as device used by this process is currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can specify device_id in init_process_group() to force use of a particular device.
Initialized distributed mode...
Namespace(config_file='configs/linea/linea_hgnetv2_l.py', options=None, coco_path='data/wireframe_processed', device='cuda', seed=42, resume='', start_epoch=0, eval=False, num_workers=2, find_unused_params=False, world_size=1, rank=0, local_rank=None, amp=False, gpu=0, distributed=True, data_aug_scales=[(1280, 1280)], data_aug_max_size=2666, data_aug_scales2_resize=[800, 1000, 1200], data_aug_scales2_crop=[768, 1200], data_aug_scale_overlap=None, batch_size_train=2, batch_size_val=4, lr=0.00025, weight_decay=0.000125, betas=[0.9, 0.999], epochs=12, lr_drop_list=[11], clip_max_norm=0.1, save_checkpoint_interval=1, modelname='LINEA', eval_spatial_size=(1280, 1280), eval_idx=5, num_classes=2, pretrained=True, use_checkpoint=False, return_interm_indices=[1, 2, 3], freeze_norm=True, freeze_stem_only=True, hybrid_encoder='hybrid_encoder_asymmetric_conv', in_channels_encoder=[512, 1024, 2048], pe_temperatureH=20, pe_temperatureW=20, transformer_activation='relu', batch_norm_type='FrozenBatchNorm2d', masks=False, aux_loss=True, num_queries=1100, query_dim=4, num_feature_levels=3, dec_n_points=[4, 1, 1], dropout=0.0, pre_norm=False, use_dn=True, dn_number=300, dn_line_noise_scale=1.0, dn_label_noise_ratio=0.5, embed_init_tgt=True, dn_labelbook_size=2, match_unstable_error=True, set_cost_class=2.0, set_cost_lines=5.0, criterionname='LINEACRITERION', criterion_type='default', weight_dict={'loss_logits': 4, 'loss_line': 5}, losses=['labels', 'lines'], focal_alpha=0.1, matcher_type='HungarianMatcher', nms_iou_threshold=-1, use_ema=False, ema_decay=0.9997, ema_epoch=0, output_dir='output/linea_hgnetv2_l', backbone='HGNetv2_B4', param_dict_type='hgnetv2_b4', use_lab=False, feat_strides=[8, 16, 32], hidden_dim=256, dim_feedforward=1024, nheads=8, use_lmap=False, expansion=0.5, depth_mult=1.0, feat_channels_decoder=[256, 256, 256], dec_layers=6, num_select=300, reg_max=16, reg_scale=4, use_warmup=False, model_parameters=[{'params': '^(?=.*backbone)(?!.*norm|bn).*$', 'lr': 1.25e-05}, {'params': '^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn)).*$', 'weight_decay': 0.0}])
Loading stage1
Loaded stage1 B4 HGNetV2 from local file.
building train_dataloader with batch_size=2...
loading annotations into memory...
Done (t=0.03s)
creating index...
index created!
building val_dataloader with batch_size=4...
loading annotations into memory...
Done (t=0.00s)
creating index...
index created!
Multi-scaling uses the following size:  [960, 992, 1024, 1056, 1088, 1120, 1152, 1184, 1216, 1248, 1280, 1280, 1280, 1312, 1344, 1376, 1408, 1440, 1472, 1504, 1536, 1568, 1600]

------------------------------------- Calculate Flops Results -------------------------------------
Notations:
number of parameters (Params), number of multiply-accumulate operations(MACs),
number of floating-point operations (FLOPs), floating-point operations per second (FLOPS),
fwd FLOPs (model forward propagation FLOPs), bwd FLOPs (model backward propagation FLOPs),
default model backpropagation takes 2.00 times as much computation as forward propagation.

Total Training Params:                                                  25.04 M 
fwd MACs:                                                               139.323 GMACs
fwd FLOPs:                                                              279.205 GFLOPS
fwd+bwd MACs:                                                           417.97 GMACs
fwd+bwd FLOPs:                                                          837.615 GFLOPS
---------------------------------------------------------------------------------------------------
{'flops': '279.205 GFLOPS', 'macs': '139.323 GMACs', 'params': 25036882}
----------------------------------------- Start training ------------------------------------------
Epoch: [0]  [  0/422]  eta: 0:25:40  lr: 0.000013  loss: 79.5078 (79.5078)  loss_line: 3.4343 (3.4343)  loss_line_0: 3.3956 (3.3956)  loss_line_1: 3.3956 (3.3956)  loss_line_2: 3.4238 (3.4238)  loss_line_3: 3.4342 (3.4342)  loss_line_4: 3.4343 (3.4343)  loss_line_dn_0: 5.7842 (5.7842)  loss_line_dn_1: 5.7907 (5.7907)  loss_line_dn_2: 5.8016 (5.8016)  loss_line_dn_3: 5.8142 (5.8142)  loss_line_dn_4: 5.8274 (5.8274)  loss_line_dn_5: 5.8407 (5.8407)  loss_line_interm: 3.5047 (3.5047)  loss_logits: 1.5700 (1.5700)  loss_logits_0: 1.4748 (1.4748)  loss_logits_1: 1.5098 (1.5098)  loss_logits_2: 1.4972 (1.4972)  loss_logits_3: 1.5098 (1.5098)  loss_logits_4: 1.5431 (1.5431)  loss_logits_dn_0: 1.6404 (1.6404)  loss_logits_dn_1: 1.6630 (1.6630)  loss_logits_dn_2: 1.7050 (1.7050)  loss_logits_dn_3: 1.7349 (1.7349)  loss_logits_dn_4: 1.7468 (1.7468)  loss_logits_dn_5: 1.7437 (1.7437)  loss_logits_interm: 1.2882 (1.2882)  time: 3.6502  data: 1.6139  max mem: 7824
Epoch: [0]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 41.1510 (44.1872)  loss_line: 0.3362 (0.7858)  loss_line_0: 0.3484 (0.8108)  loss_line_1: 0.3302 (0.7915)  loss_line_2: 0.3381 (0.7889)  loss_line_3: 0.3425 (0.7865)  loss_line_4: 0.3370 (0.7859)  loss_line_dn_0: 4.9401 (5.0422)  loss_line_dn_1: 4.9372 (5.0353)  loss_line_dn_2: 4.9379 (5.0338)  loss_line_dn_3: 4.9430 (5.0330)  loss_line_dn_4: 4.9441 (5.0325)  loss_line_dn_5: 4.9442 (5.0324)  loss_line_interm: 0.3565 (0.7861)  loss_logits: 0.7384 (0.8784)  loss_logits_0: 0.7873 (0.8860)  loss_logits_1: 0.7798 (0.8810)  loss_logits_2: 0.7487 (0.8852)  loss_logits_3: 0.7418 (0.8851)  loss_logits_4: 0.7477 (0.8752)  loss_logits_dn_0: 0.3304 (0.3603)  loss_logits_dn_1: 0.3274 (0.3623)  loss_logits_dn_2: 0.3296 (0.3653)  loss_logits_dn_3: 0.3278 (0.3702)  loss_logits_dn_4: 0.3302 (0.3789)  loss_logits_dn_5: 0.3321 (0.3944)  loss_logits_interm: 0.9200 (0.9201)  time: 1.1838  data: 0.0079  max mem: 11995
Epoch: [0] Total time: 0:08:09 (1.1604 s / it)
Averaged stats: lr: 0.000013  loss: 41.1510 (44.1872)  loss_line: 0.3362 (0.7858)  loss_line_0: 0.3484 (0.8108)  loss_line_1: 0.3302 (0.7915)  loss_line_2: 0.3381 (0.7889)  loss_line_3: 0.3425 (0.7865)  loss_line_4: 0.3370 (0.7859)  loss_line_dn_0: 4.9401 (5.0422)  loss_line_dn_1: 4.9372 (5.0353)  loss_line_dn_2: 4.9379 (5.0338)  loss_line_dn_3: 4.9430 (5.0330)  loss_line_dn_4: 4.9441 (5.0325)  loss_line_dn_5: 4.9442 (5.0324)  loss_line_interm: 0.3565 (0.7861)  loss_logits: 0.7384 (0.8784)  loss_logits_0: 0.7873 (0.8860)  loss_logits_1: 0.7798 (0.8810)  loss_logits_2: 0.7487 (0.8852)  loss_logits_3: 0.7418 (0.8851)  loss_logits_4: 0.7477 (0.8752)  loss_logits_dn_0: 0.3304 (0.3603)  loss_logits_dn_1: 0.3274 (0.3623)  loss_logits_dn_2: 0.3296 (0.3653)  loss_logits_dn_3: 0.3278 (0.3702)  loss_logits_dn_4: 0.3302 (0.3789)  loss_logits_dn_5: 0.3321 (0.3944)  loss_logits_interm: 0.9200 (0.9201)
Test:  [ 0/25]  eta: 0:00:53  loss: 1.0384 (1.0384)  loss_line: 0.3253 (0.3253)  loss_logits: 0.7131 (0.7131)  time: 2.1272  data: 1.3862  max mem: 11995
Test:  [24/25]  eta: 0:00:00  loss: 1.0410 (1.0733)  loss_line: 0.2904 (0.3079)  loss_logits: 0.7280 (0.7655)  time: 0.6506  data: 0.0490  max mem: 11995
Test: Total time: 0:00:17 (0.7139 s / it)
Averaged stats: loss: 1.0410 (1.0733)  loss_line: 0.2904 (0.3079)  loss_logits: 0.7280 (0.7655)
Epoch: [1]  [  0/422]  eta: 0:12:22  lr: 0.000013  loss: 40.1220 (40.1220)  loss_line: 0.6914 (0.6914)  loss_line_0: 0.6746 (0.6746)  loss_line_1: 0.6966 (0.6966)  loss_line_2: 0.6947 (0.6947)  loss_line_3: 0.6896 (0.6896)  loss_line_4: 0.6912 (0.6912)  loss_line_dn_0: 4.7286 (4.7286)  loss_line_dn_1: 4.7360 (4.7360)  loss_line_dn_2: 4.7894 (4.7894)  loss_line_dn_3: 4.7430 (4.7430)  loss_line_dn_4: 4.7238 (4.7238)  loss_line_dn_5: 4.7184 (4.7184)  loss_line_interm: 0.6652 (0.6652)  loss_logits: 0.6713 (0.6713)  loss_logits_0: 0.7156 (0.7156)  loss_logits_1: 0.6854 (0.6854)  loss_logits_2: 0.6499 (0.6499)  loss_logits_3: 0.6826 (0.6826)  loss_logits_4: 0.6992 (0.6992)  loss_logits_dn_0: 0.3338 (0.3338)  loss_logits_dn_1: 0.3290 (0.3290)  loss_logits_dn_2: 0.3280 (0.3280)  loss_logits_dn_3: 0.3319 (0.3319)  loss_logits_dn_4: 0.3360 (0.3360)  loss_logits_dn_5: 0.3357 (0.3357)  loss_logits_interm: 0.7813 (0.7813)  time: 1.7595  data: 0.8841  max mem: 11995
Epoch: [1]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 40.1447 (41.1630)  loss_line: 0.4723 (0.4601)  loss_line_0: 0.4804 (0.4620)  loss_line_1: 0.4909 (0.4601)  loss_line_2: 0.4912 (0.4610)  loss_line_3: 0.4850 (0.4609)  loss_line_4: 0.4721 (0.4607)  loss_line_dn_0: 4.9563 (5.1041)  loss_line_dn_1: 5.0415 (5.1075)  loss_line_dn_2: 5.0425 (5.1089)  loss_line_dn_3: 5.0432 (5.1093)  loss_line_dn_4: 5.0434 (5.1096)  loss_line_dn_5: 5.0436 (5.1098)  loss_line_interm: 0.4960 (0.4592)  loss_logits: 0.6270 (0.7261)  loss_logits_0: 0.7382 (0.8076)  loss_logits_1: 0.7269 (0.7632)  loss_logits_2: 0.6747 (0.7381)  loss_logits_3: 0.6568 (0.7350)  loss_logits_4: 0.6418 (0.7316)  loss_logits_dn_0: 0.3276 (0.3224)  loss_logits_dn_1: 0.3269 (0.3191)  loss_logits_dn_2: 0.3274 (0.3203)  loss_logits_dn_3: 0.3280 (0.3220)  loss_logits_dn_4: 0.3290 (0.3245)  loss_logits_dn_5: 0.3314 (0.3283)  loss_logits_interm: 0.7810 (0.8516)  time: 1.0684  data: 0.0087  max mem: 11995
Epoch: [1] Total time: 0:08:05 (1.1494 s / it)
Averaged stats: lr: 0.000013  loss: 40.1447 (41.1630)  loss_line: 0.4723 (0.4601)  loss_line_0: 0.4804 (0.4620)  loss_line_1: 0.4909 (0.4601)  loss_line_2: 0.4912 (0.4610)  loss_line_3: 0.4850 (0.4609)  loss_line_4: 0.4721 (0.4607)  loss_line_dn_0: 4.9563 (5.1041)  loss_line_dn_1: 5.0415 (5.1075)  loss_line_dn_2: 5.0425 (5.1089)  loss_line_dn_3: 5.0432 (5.1093)  loss_line_dn_4: 5.0434 (5.1096)  loss_line_dn_5: 5.0436 (5.1098)  loss_line_interm: 0.4960 (0.4592)  loss_logits: 0.6270 (0.7261)  loss_logits_0: 0.7382 (0.8076)  loss_logits_1: 0.7269 (0.7632)  loss_logits_2: 0.6747 (0.7381)  loss_logits_3: 0.6568 (0.7350)  loss_logits_4: 0.6418 (0.7316)  loss_logits_dn_0: 0.3276 (0.3224)  loss_logits_dn_1: 0.3269 (0.3191)  loss_logits_dn_2: 0.3274 (0.3203)  loss_logits_dn_3: 0.3280 (0.3220)  loss_logits_dn_4: 0.3290 (0.3245)  loss_logits_dn_5: 0.3314 (0.3283)  loss_logits_interm: 0.7810 (0.8516)
Test:  [ 0/25]  eta: 0:00:45  loss: 1.0232 (1.0232)  loss_line: 0.2803 (0.2803)  loss_logits: 0.7429 (0.7429)  time: 1.8329  data: 1.0999  max mem: 11995
Test:  [24/25]  eta: 0:00:00  loss: 0.9801 (1.0079)  loss_line: 0.2912 (0.3049)  loss_logits: 0.6888 (0.7030)  time: 0.6425  data: 0.0442  max mem: 11995
Test: Total time: 0:00:18 (0.7243 s / it)
Averaged stats: loss: 0.9801 (1.0079)  loss_line: 0.2912 (0.3049)  loss_logits: 0.6888 (0.7030)
Epoch: [2]  [  0/422]  eta: 0:15:14  lr: 0.000013  loss: 39.0184 (39.0184)  loss_line: 0.2325 (0.2325)  loss_line_0: 0.2010 (0.2010)  loss_line_1: 0.2018 (0.2018)  loss_line_2: 0.2316 (0.2316)  loss_line_3: 0.2085 (0.2085)  loss_line_4: 0.2242 (0.2242)  loss_line_dn_0: 4.8578 (4.8578)  loss_line_dn_1: 4.8473 (4.8473)  loss_line_dn_2: 4.8440 (4.8440)  loss_line_dn_3: 4.8441 (4.8441)  loss_line_dn_4: 4.8429 (4.8429)  loss_line_dn_5: 4.8424 (4.8424)  loss_line_interm: 0.2056 (0.2056)  loss_logits: 0.9411 (0.9411)  loss_logits_0: 1.0549 (1.0549)  loss_logits_1: 0.9779 (0.9779)  loss_logits_2: 0.8656 (0.8656)  loss_logits_3: 0.8980 (0.8980)  loss_logits_4: 0.8641 (0.8641)  loss_logits_dn_0: 0.2784 (0.2784)  loss_logits_dn_1: 0.2772 (0.2772)  loss_logits_dn_2: 0.2995 (0.2995)  loss_logits_dn_3: 0.3112 (0.3112)  loss_logits_dn_4: 0.3015 (0.3015)  loss_logits_dn_5: 0.3134 (0.3134)  loss_logits_interm: 1.0519 (1.0519)  time: 2.1663  data: 0.9997  max mem: 11995
Epoch: [2]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 37.0222 (40.3988)  loss_line: 0.2733 (0.4635)  loss_line_0: 0.2859 (0.4696)  loss_line_1: 0.2935 (0.4660)  loss_line_2: 0.2716 (0.4643)  loss_line_3: 0.2738 (0.4634)  loss_line_4: 0.2810 (0.4634)  loss_line_dn_0: 4.7753 (5.0300)  loss_line_dn_1: 4.7742 (5.0174)  loss_line_dn_2: 4.7737 (5.0161)  loss_line_dn_3: 4.7777 (5.0157)  loss_line_dn_4: 4.7823 (5.0161)  loss_line_dn_5: 4.7828 (5.0164)  loss_line_interm: 0.2856 (0.4788)  loss_logits: 0.6288 (0.6893)  loss_logits_0: 0.6822 (0.7733)  loss_logits_1: 0.6337 (0.7279)  loss_logits_2: 0.6488 (0.7116)  loss_logits_3: 0.6157 (0.7070)  loss_logits_4: 0.6339 (0.6953)  loss_logits_dn_0: 0.3226 (0.3207)  loss_logits_dn_1: 0.3238 (0.3135)  loss_logits_dn_2: 0.3210 (0.3140)  loss_logits_dn_3: 0.3207 (0.3159)  loss_logits_dn_4: 0.3250 (0.3193)  loss_logits_dn_5: 0.3256 (0.3233)  loss_logits_interm: 0.7278 (0.8071)  time: 1.0999  data: 0.0093  max mem: 11995
Epoch: [2] Total time: 0:08:08 (1.1570 s / it)
Averaged stats: lr: 0.000013  loss: 37.0222 (40.3988)  loss_line: 0.2733 (0.4635)  loss_line_0: 0.2859 (0.4696)  loss_line_1: 0.2935 (0.4660)  loss_line_2: 0.2716 (0.4643)  loss_line_3: 0.2738 (0.4634)  loss_line_4: 0.2810 (0.4634)  loss_line_dn_0: 4.7753 (5.0300)  loss_line_dn_1: 4.7742 (5.0174)  loss_line_dn_2: 4.7737 (5.0161)  loss_line_dn_3: 4.7777 (5.0157)  loss_line_dn_4: 4.7823 (5.0161)  loss_line_dn_5: 4.7828 (5.0164)  loss_line_interm: 0.2856 (0.4788)  loss_logits: 0.6288 (0.6893)  loss_logits_0: 0.6822 (0.7733)  loss_logits_1: 0.6337 (0.7279)  loss_logits_2: 0.6488 (0.7116)  loss_logits_3: 0.6157 (0.7070)  loss_logits_4: 0.6339 (0.6953)  loss_logits_dn_0: 0.3226 (0.3207)  loss_logits_dn_1: 0.3238 (0.3135)  loss_logits_dn_2: 0.3210 (0.3140)  loss_logits_dn_3: 0.3207 (0.3159)  loss_logits_dn_4: 0.3250 (0.3193)  loss_logits_dn_5: 0.3256 (0.3233)  loss_logits_interm: 0.7278 (0.8071)
Test:  [ 0/25]  eta: 0:01:01  loss: 0.8453 (0.8453)  loss_line: 0.3205 (0.3205)  loss_logits: 0.5248 (0.5248)  time: 2.4460  data: 1.7065  max mem: 11995
Test:  [24/25]  eta: 0:00:00  loss: 0.8545 (0.8978)  loss_line: 0.2557 (0.2577)  loss_logits: 0.6141 (0.6401)  time: 0.6368  data: 0.0408  max mem: 11995
Test: Total time: 0:00:18 (0.7215 s / it)
Averaged stats: loss: 0.8545 (0.8978)  loss_line: 0.2557 (0.2577)  loss_logits: 0.6141 (0.6401)
Epoch: [3]  [  0/422]  eta: 0:13:10  lr: 0.000013  loss: 44.2360 (44.2360)  loss_line: 0.1453 (0.1453)  loss_line_0: 0.1549 (0.1549)  loss_line_1: 0.1476 (0.1476)  loss_line_2: 0.1365 (0.1365)  loss_line_3: 0.1396 (0.1396)  loss_line_4: 0.1429 (0.1429)  loss_line_dn_0: 6.1966 (6.1966)  loss_line_dn_1: 6.1997 (6.1997)  loss_line_dn_2: 6.2023 (6.2023)  loss_line_dn_3: 6.2025 (6.2025)  loss_line_dn_4: 6.2028 (6.2028)  loss_line_dn_5: 6.2032 (6.2032)  loss_line_interm: 0.2073 (0.2073)  loss_logits: 0.5545 (0.5545)  loss_logits_0: 0.5683 (0.5683)  loss_logits_1: 0.5724 (0.5724)  loss_logits_2: 0.6111 (0.6111)  loss_logits_3: 0.5571 (0.5571)  loss_logits_4: 0.5496 (0.5496)  loss_logits_dn_0: 0.3199 (0.3199)  loss_logits_dn_1: 0.3223 (0.3223)  loss_logits_dn_2: 0.3198 (0.3198)  loss_logits_dn_3: 0.3202 (0.3202)  loss_logits_dn_4: 0.3213 (0.3213)  loss_logits_dn_5: 0.3220 (0.3220)  loss_logits_interm: 0.6164 (0.6164)  time: 1.8731  data: 0.9287  max mem: 11995
Epoch: [3]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 36.1374 (39.7223)  loss_line: 0.2571 (0.4501)  loss_line_0: 0.3161 (0.4773)  loss_line_1: 0.2772 (0.4654)  loss_line_2: 0.2750 (0.4556)  loss_line_3: 0.2839 (0.4551)  loss_line_4: 0.2823 (0.4531)  loss_line_dn_0: 4.5778 (5.0125)  loss_line_dn_1: 4.5587 (5.0034)  loss_line_dn_2: 4.5577 (5.0017)  loss_line_dn_3: 4.5580 (5.0015)  loss_line_dn_4: 4.5567 (5.0017)  loss_line_dn_5: 4.5560 (5.0018)  loss_line_interm: 0.3641 (0.4894)  loss_logits: 0.6079 (0.6353)  loss_logits_0: 0.5916 (0.6765)  loss_logits_1: 0.6167 (0.6568)  loss_logits_2: 0.5951 (0.6401)  loss_logits_3: 0.5825 (0.6332)  loss_logits_4: 0.6082 (0.6351)  loss_logits_dn_0: 0.3168 (0.3150)  loss_logits_dn_1: 0.3129 (0.3035)  loss_logits_dn_2: 0.3132 (0.3018)  loss_logits_dn_3: 0.3108 (0.3038)  loss_logits_dn_4: 0.3130 (0.3067)  loss_logits_dn_5: 0.3228 (0.3119)  loss_logits_interm: 0.6744 (0.7344)  time: 1.1312  data: 0.0106  max mem: 11996
Epoch: [3] Total time: 0:08:24 (1.1962 s / it)
Averaged stats: lr: 0.000013  loss: 36.1374 (39.7223)  loss_line: 0.2571 (0.4501)  loss_line_0: 0.3161 (0.4773)  loss_line_1: 0.2772 (0.4654)  loss_line_2: 0.2750 (0.4556)  loss_line_3: 0.2839 (0.4551)  loss_line_4: 0.2823 (0.4531)  loss_line_dn_0: 4.5778 (5.0125)  loss_line_dn_1: 4.5587 (5.0034)  loss_line_dn_2: 4.5577 (5.0017)  loss_line_dn_3: 4.5580 (5.0015)  loss_line_dn_4: 4.5567 (5.0017)  loss_line_dn_5: 4.5560 (5.0018)  loss_line_interm: 0.3641 (0.4894)  loss_logits: 0.6079 (0.6353)  loss_logits_0: 0.5916 (0.6765)  loss_logits_1: 0.6167 (0.6568)  loss_logits_2: 0.5951 (0.6401)  loss_logits_3: 0.5825 (0.6332)  loss_logits_4: 0.6082 (0.6351)  loss_logits_dn_0: 0.3168 (0.3150)  loss_logits_dn_1: 0.3129 (0.3035)  loss_logits_dn_2: 0.3132 (0.3018)  loss_logits_dn_3: 0.3108 (0.3038)  loss_logits_dn_4: 0.3130 (0.3067)  loss_logits_dn_5: 0.3228 (0.3119)  loss_logits_interm: 0.6744 (0.7344)
Test:  [ 0/25]  eta: 0:00:43  loss: 0.8490 (0.8490)  loss_line: 0.2310 (0.2310)  loss_logits: 0.6179 (0.6179)  time: 1.7234  data: 1.0148  max mem: 11996
Test:  [24/25]  eta: 0:00:00  loss: 0.8859 (0.8897)  loss_line: 0.2354 (0.2292)  loss_logits: 0.6240 (0.6605)  time: 0.6526  data: 0.0555  max mem: 11996
Test: Total time: 0:00:17 (0.7016 s / it)
Averaged stats: loss: 0.8859 (0.8897)  loss_line: 0.2354 (0.2292)  loss_logits: 0.6240 (0.6605)
Epoch: [4]  [  0/422]  eta: 0:18:02  lr: 0.000013  loss: 33.6893 (33.6893)  loss_line: 0.1568 (0.1568)  loss_line_0: 0.1469 (0.1469)  loss_line_1: 0.1431 (0.1431)  loss_line_2: 0.1834 (0.1834)  loss_line_3: 0.1349 (0.1349)  loss_line_4: 0.1697 (0.1697)  loss_line_dn_0: 4.1770 (4.1770)  loss_line_dn_1: 4.1721 (4.1721)  loss_line_dn_2: 4.1693 (4.1693)  loss_line_dn_3: 4.1663 (4.1663)  loss_line_dn_4: 4.1657 (4.1657)  loss_line_dn_5: 4.1653 (4.1653)  loss_line_interm: 0.1578 (0.1578)  loss_logits: 0.8997 (0.8997)  loss_logits_0: 0.7755 (0.7755)  loss_logits_1: 0.8014 (0.8014)  loss_logits_2: 0.7666 (0.7666)  loss_logits_3: 0.7897 (0.7897)  loss_logits_4: 0.8534 (0.8534)  loss_logits_dn_0: 0.3054 (0.3054)  loss_logits_dn_1: 0.3046 (0.3046)  loss_logits_dn_2: 0.2940 (0.2940)  loss_logits_dn_3: 0.2953 (0.2953)  loss_logits_dn_4: 0.2874 (0.2874)  loss_logits_dn_5: 0.2935 (0.2935)  loss_logits_interm: 0.9145 (0.9145)  time: 2.5650  data: 1.0553  max mem: 11996
Epoch: [4]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 38.0780 (38.8373)  loss_line: 0.2261 (0.3694)  loss_line_0: 0.2620 (0.4022)  loss_line_1: 0.1953 (0.3896)  loss_line_2: 0.1955 (0.3797)  loss_line_3: 0.2246 (0.3753)  loss_line_4: 0.2276 (0.3720)  loss_line_dn_0: 4.7010 (4.9632)  loss_line_dn_1: 4.7136 (4.9544)  loss_line_dn_2: 4.7214 (4.9524)  loss_line_dn_3: 4.7229 (4.9518)  loss_line_dn_4: 4.7243 (4.9518)  loss_line_dn_5: 4.7253 (4.9519)  loss_line_interm: 0.2670 (0.4286)  loss_logits: 0.6086 (0.6273)  loss_logits_0: 0.6924 (0.6723)  loss_logits_1: 0.6510 (0.6507)  loss_logits_2: 0.6111 (0.6380)  loss_logits_3: 0.6039 (0.6314)  loss_logits_4: 0.5990 (0.6271)  loss_logits_dn_0: 0.3063 (0.3110)  loss_logits_dn_1: 0.3029 (0.2985)  loss_logits_dn_2: 0.2936 (0.2965)  loss_logits_dn_3: 0.2943 (0.2976)  loss_logits_dn_4: 0.2992 (0.2997)  loss_logits_dn_5: 0.3013 (0.3028)  loss_logits_interm: 0.7456 (0.7420)  time: 1.1575  data: 0.0095  max mem: 11996
Epoch: [4] Total time: 0:08:19 (1.1830 s / it)
Averaged stats: lr: 0.000013  loss: 38.0780 (38.8373)  loss_line: 0.2261 (0.3694)  loss_line_0: 0.2620 (0.4022)  loss_line_1: 0.1953 (0.3896)  loss_line_2: 0.1955 (0.3797)  loss_line_3: 0.2246 (0.3753)  loss_line_4: 0.2276 (0.3720)  loss_line_dn_0: 4.7010 (4.9632)  loss_line_dn_1: 4.7136 (4.9544)  loss_line_dn_2: 4.7214 (4.9524)  loss_line_dn_3: 4.7229 (4.9518)  loss_line_dn_4: 4.7243 (4.9518)  loss_line_dn_5: 4.7253 (4.9519)  loss_line_interm: 0.2670 (0.4286)  loss_logits: 0.6086 (0.6273)  loss_logits_0: 0.6924 (0.6723)  loss_logits_1: 0.6510 (0.6507)  loss_logits_2: 0.6111 (0.6380)  loss_logits_3: 0.6039 (0.6314)  loss_logits_4: 0.5990 (0.6271)  loss_logits_dn_0: 0.3063 (0.3110)  loss_logits_dn_1: 0.3029 (0.2985)  loss_logits_dn_2: 0.2936 (0.2965)  loss_logits_dn_3: 0.2943 (0.2976)  loss_logits_dn_4: 0.2992 (0.2997)  loss_logits_dn_5: 0.3013 (0.3028)  loss_logits_interm: 0.7456 (0.7420)
Test:  [ 0/25]  eta: 0:00:58  loss: 0.6581 (0.6581)  loss_line: 0.1240 (0.1240)  loss_logits: 0.5341 (0.5341)  time: 2.3420  data: 1.5323  max mem: 11996
Test:  [24/25]  eta: 0:00:00  loss: 0.7860 (0.8269)  loss_line: 0.1837 (0.1784)  loss_logits: 0.6169 (0.6485)  time: 0.6510  data: 0.0479  max mem: 11996
Test: Total time: 0:00:18 (0.7271 s / it)
Averaged stats: loss: 0.7860 (0.8269)  loss_line: 0.1837 (0.1784)  loss_logits: 0.6169 (0.6485)
Epoch: [5]  [  0/422]  eta: 0:19:32  lr: 0.000013  loss: 36.6659 (36.6659)  loss_line: 0.1949 (0.1949)  loss_line_0: 0.2381 (0.2381)  loss_line_1: 0.2313 (0.2313)  loss_line_2: 0.2067 (0.2067)  loss_line_3: 0.2072 (0.2072)  loss_line_4: 0.1957 (0.1957)  loss_line_dn_0: 4.7505 (4.7505)  loss_line_dn_1: 4.7425 (4.7425)  loss_line_dn_2: 4.7382 (4.7382)  loss_line_dn_3: 4.7377 (4.7377)  loss_line_dn_4: 4.7371 (4.7371)  loss_line_dn_5: 4.7374 (4.7374)  loss_line_interm: 0.2575 (0.2575)  loss_logits: 0.6656 (0.6656)  loss_logits_0: 0.7414 (0.7414)  loss_logits_1: 0.6464 (0.6464)  loss_logits_2: 0.6566 (0.6566)  loss_logits_3: 0.7066 (0.7066)  loss_logits_4: 0.6866 (0.6866)  loss_logits_dn_0: 0.3067 (0.3067)  loss_logits_dn_1: 0.3055 (0.3055)  loss_logits_dn_2: 0.3027 (0.3027)  loss_logits_dn_3: 0.3046 (0.3046)  loss_logits_dn_4: 0.3019 (0.3019)  loss_logits_dn_5: 0.3043 (0.3043)  loss_logits_interm: 0.7622 (0.7622)  time: 2.7777  data: 1.1110  max mem: 11996
Epoch: [5]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 42.1864 (38.8240)  loss_line: 0.1979 (0.3721)  loss_line_0: 0.2364 (0.4044)  loss_line_1: 0.2178 (0.3925)  loss_line_2: 0.2013 (0.3794)  loss_line_3: 0.1926 (0.3765)  loss_line_4: 0.1983 (0.3745)  loss_line_dn_0: 5.3656 (5.0089)  loss_line_dn_1: 5.3846 (4.9981)  loss_line_dn_2: 5.3983 (4.9930)  loss_line_dn_3: 5.3973 (4.9916)  loss_line_dn_4: 5.3957 (4.9914)  loss_line_dn_5: 5.3922 (4.9914)  loss_line_interm: 0.2618 (0.4385)  loss_logits: 0.5147 (0.5824)  loss_logits_0: 0.6454 (0.6431)  loss_logits_1: 0.6217 (0.6172)  loss_logits_2: 0.5547 (0.6104)  loss_logits_3: 0.5362 (0.5982)  loss_logits_4: 0.5541 (0.5904)  loss_logits_dn_0: 0.3058 (0.3090)  loss_logits_dn_1: 0.2942 (0.2949)  loss_logits_dn_2: 0.2950 (0.2923)  loss_logits_dn_3: 0.2948 (0.2922)  loss_logits_dn_4: 0.2933 (0.2926)  loss_logits_dn_5: 0.2937 (0.2933)  loss_logits_interm: 0.6524 (0.6955)  time: 1.1320  data: 0.0093  max mem: 11996
Epoch: [5] Total time: 0:08:10 (1.1613 s / it)
Averaged stats: lr: 0.000013  loss: 42.1864 (38.8240)  loss_line: 0.1979 (0.3721)  loss_line_0: 0.2364 (0.4044)  loss_line_1: 0.2178 (0.3925)  loss_line_2: 0.2013 (0.3794)  loss_line_3: 0.1926 (0.3765)  loss_line_4: 0.1983 (0.3745)  loss_line_dn_0: 5.3656 (5.0089)  loss_line_dn_1: 5.3846 (4.9981)  loss_line_dn_2: 5.3983 (4.9930)  loss_line_dn_3: 5.3973 (4.9916)  loss_line_dn_4: 5.3957 (4.9914)  loss_line_dn_5: 5.3922 (4.9914)  loss_line_interm: 0.2618 (0.4385)  loss_logits: 0.5147 (0.5824)  loss_logits_0: 0.6454 (0.6431)  loss_logits_1: 0.6217 (0.6172)  loss_logits_2: 0.5547 (0.6104)  loss_logits_3: 0.5362 (0.5982)  loss_logits_4: 0.5541 (0.5904)  loss_logits_dn_0: 0.3058 (0.3090)  loss_logits_dn_1: 0.2942 (0.2949)  loss_logits_dn_2: 0.2950 (0.2923)  loss_logits_dn_3: 0.2948 (0.2922)  loss_logits_dn_4: 0.2933 (0.2926)  loss_logits_dn_5: 0.2937 (0.2933)  loss_logits_interm: 0.6524 (0.6955)
Test:  [ 0/25]  eta: 0:00:44  loss: 0.7166 (0.7166)  loss_line: 0.1268 (0.1268)  loss_logits: 0.5897 (0.5897)  time: 1.7720  data: 1.0239  max mem: 11996
Test:  [24/25]  eta: 0:00:00  loss: 0.7251 (0.7800)  loss_line: 0.1462 (0.1537)  loss_logits: 0.5806 (0.6263)  time: 0.6452  data: 0.0445  max mem: 11996
Test: Total time: 0:00:17 (0.7003 s / it)
Averaged stats: loss: 0.7251 (0.7800)  loss_line: 0.1462 (0.1537)  loss_logits: 0.5806 (0.6263)
Epoch: [6]  [  0/422]  eta: 0:19:51  lr: 0.000013  loss: 43.0455 (43.0455)  loss_line: 1.2213 (1.2213)  loss_line_0: 1.3305 (1.3305)  loss_line_1: 1.3267 (1.3267)  loss_line_2: 1.2795 (1.2795)  loss_line_3: 1.2394 (1.2394)  loss_line_4: 1.2537 (1.2537)  loss_line_dn_0: 4.4934 (4.4934)  loss_line_dn_1: 4.4776 (4.4776)  loss_line_dn_2: 4.4835 (4.4835)  loss_line_dn_3: 4.4857 (4.4857)  loss_line_dn_4: 4.4856 (4.4856)  loss_line_dn_5: 4.4841 (4.4841)  loss_line_interm: 1.3451 (1.3451)  loss_logits: 1.1284 (1.1284)  loss_logits_0: 0.7939 (0.7939)  loss_logits_1: 0.6577 (0.6577)  loss_logits_2: 0.7447 (0.7447)  loss_logits_3: 0.8477 (0.8477)  loss_logits_4: 0.8158 (0.8158)  loss_logits_dn_0: 0.2471 (0.2471)  loss_logits_dn_1: 0.2041 (0.2041)  loss_logits_dn_2: 0.1790 (0.1790)  loss_logits_dn_3: 0.1688 (0.1688)  loss_logits_dn_4: 0.1673 (0.1673)  loss_logits_dn_5: 0.1636 (0.1636)  loss_logits_interm: 1.0213 (1.0213)  time: 2.8231  data: 1.5511  max mem: 11996
Epoch: [6]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 38.0630 (38.2972)  loss_line: 0.1448 (0.3341)  loss_line_0: 0.1875 (0.3740)  loss_line_1: 0.1600 (0.3640)  loss_line_2: 0.1484 (0.3453)  loss_line_3: 0.1558 (0.3413)  loss_line_4: 0.1669 (0.3355)  loss_line_dn_0: 5.1571 (4.9525)  loss_line_dn_1: 5.1535 (4.9398)  loss_line_dn_2: 5.1296 (4.9366)  loss_line_dn_3: 5.1307 (4.9356)  loss_line_dn_4: 5.1242 (4.9350)  loss_line_dn_5: 5.1229 (4.9350)  loss_line_interm: 0.2298 (0.4065)  loss_logits: 0.6194 (0.5970)  loss_logits_0: 0.6487 (0.6466)  loss_logits_1: 0.6404 (0.6364)  loss_logits_2: 0.6187 (0.6279)  loss_logits_3: 0.5782 (0.5986)  loss_logits_4: 0.5916 (0.5977)  loss_logits_dn_0: 0.2847 (0.3050)  loss_logits_dn_1: 0.2739 (0.2922)  loss_logits_dn_2: 0.2677 (0.2900)  loss_logits_dn_3: 0.2698 (0.2891)  loss_logits_dn_4: 0.2723 (0.2904)  loss_logits_dn_5: 0.2734 (0.2937)  loss_logits_interm: 0.7125 (0.6976)  time: 1.1731  data: 0.0105  max mem: 11997
Epoch: [6] Total time: 0:08:23 (1.1922 s / it)
Averaged stats: lr: 0.000013  loss: 38.0630 (38.2972)  loss_line: 0.1448 (0.3341)  loss_line_0: 0.1875 (0.3740)  loss_line_1: 0.1600 (0.3640)  loss_line_2: 0.1484 (0.3453)  loss_line_3: 0.1558 (0.3413)  loss_line_4: 0.1669 (0.3355)  loss_line_dn_0: 5.1571 (4.9525)  loss_line_dn_1: 5.1535 (4.9398)  loss_line_dn_2: 5.1296 (4.9366)  loss_line_dn_3: 5.1307 (4.9356)  loss_line_dn_4: 5.1242 (4.9350)  loss_line_dn_5: 5.1229 (4.9350)  loss_line_interm: 0.2298 (0.4065)  loss_logits: 0.6194 (0.5970)  loss_logits_0: 0.6487 (0.6466)  loss_logits_1: 0.6404 (0.6364)  loss_logits_2: 0.6187 (0.6279)  loss_logits_3: 0.5782 (0.5986)  loss_logits_4: 0.5916 (0.5977)  loss_logits_dn_0: 0.2847 (0.3050)  loss_logits_dn_1: 0.2739 (0.2922)  loss_logits_dn_2: 0.2677 (0.2900)  loss_logits_dn_3: 0.2698 (0.2891)  loss_logits_dn_4: 0.2723 (0.2904)  loss_logits_dn_5: 0.2734 (0.2937)  loss_logits_interm: 0.7125 (0.6976)
Test:  [ 0/25]  eta: 0:00:44  loss: 0.5968 (0.5968)  loss_line: 0.1173 (0.1173)  loss_logits: 0.4795 (0.4795)  time: 1.7801  data: 1.0947  max mem: 11997
Test:  [24/25]  eta: 0:00:00  loss: 0.6787 (0.7028)  loss_line: 0.1404 (0.1376)  loss_logits: 0.5291 (0.5651)  time: 0.6431  data: 0.0407  max mem: 11997
Test: Total time: 0:00:17 (0.6960 s / it)
Averaged stats: loss: 0.6787 (0.7028)  loss_line: 0.1404 (0.1376)  loss_logits: 0.5291 (0.5651)
Epoch: [7]  [  0/422]  eta: 0:19:37  lr: 0.000013  loss: 28.8348 (28.8348)  loss_line: 0.1959 (0.1959)  loss_line_0: 0.2276 (0.2276)  loss_line_1: 0.2302 (0.2302)  loss_line_2: 0.2067 (0.2067)  loss_line_3: 0.2100 (0.2100)  loss_line_4: 0.1983 (0.1983)  loss_line_dn_0: 3.5207 (3.5207)  loss_line_dn_1: 3.5122 (3.5122)  loss_line_dn_2: 3.5082 (3.5082)  loss_line_dn_3: 3.5064 (3.5064)  loss_line_dn_4: 3.5058 (3.5058)  loss_line_dn_5: 3.5062 (3.5062)  loss_line_interm: 0.2242 (0.2242)  loss_logits: 0.6064 (0.6064)  loss_logits_0: 0.6827 (0.6827)  loss_logits_1: 0.6274 (0.6274)  loss_logits_2: 0.6399 (0.6399)  loss_logits_3: 0.6134 (0.6134)  loss_logits_4: 0.5721 (0.5721)  loss_logits_dn_0: 0.3114 (0.3114)  loss_logits_dn_1: 0.2927 (0.2927)  loss_logits_dn_2: 0.2929 (0.2929)  loss_logits_dn_3: 0.2898 (0.2898)  loss_logits_dn_4: 0.2929 (0.2929)  loss_logits_dn_5: 0.2958 (0.2958)  loss_logits_interm: 0.7651 (0.7651)  time: 2.7892  data: 1.1309  max mem: 11997
Epoch: [7]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 37.1293 (38.1229)  loss_line: 0.1708 (0.3578)  loss_line_0: 0.2116 (0.3906)  loss_line_1: 0.2338 (0.3822)  loss_line_2: 0.1954 (0.3667)  loss_line_3: 0.1862 (0.3613)  loss_line_4: 0.1857 (0.3597)  loss_line_dn_0: 4.9051 (4.9504)  loss_line_dn_1: 4.8754 (4.9364)  loss_line_dn_2: 4.8714 (4.9340)  loss_line_dn_3: 4.8707 (4.9332)  loss_line_dn_4: 4.8706 (4.9328)  loss_line_dn_5: 4.8707 (4.9329)  loss_line_interm: 0.2662 (0.4257)  loss_logits: 0.5142 (0.5599)  loss_logits_0: 0.5897 (0.5967)  loss_logits_1: 0.5541 (0.5886)  loss_logits_2: 0.5363 (0.5809)  loss_logits_3: 0.5369 (0.5666)  loss_logits_4: 0.5127 (0.5601)  loss_logits_dn_0: 0.3044 (0.3039)  loss_logits_dn_1: 0.2849 (0.2905)  loss_logits_dn_2: 0.2869 (0.2873)  loss_logits_dn_3: 0.2857 (0.2866)  loss_logits_dn_4: 0.2924 (0.2872)  loss_logits_dn_5: 0.2970 (0.2885)  loss_logits_interm: 0.7093 (0.6627)  time: 1.1275  data: 0.0091  max mem: 11997
Epoch: [7] Total time: 0:08:11 (1.1643 s / it)
Averaged stats: lr: 0.000013  loss: 37.1293 (38.1229)  loss_line: 0.1708 (0.3578)  loss_line_0: 0.2116 (0.3906)  loss_line_1: 0.2338 (0.3822)  loss_line_2: 0.1954 (0.3667)  loss_line_3: 0.1862 (0.3613)  loss_line_4: 0.1857 (0.3597)  loss_line_dn_0: 4.9051 (4.9504)  loss_line_dn_1: 4.8754 (4.9364)  loss_line_dn_2: 4.8714 (4.9340)  loss_line_dn_3: 4.8707 (4.9332)  loss_line_dn_4: 4.8706 (4.9328)  loss_line_dn_5: 4.8707 (4.9329)  loss_line_interm: 0.2662 (0.4257)  loss_logits: 0.5142 (0.5599)  loss_logits_0: 0.5897 (0.5967)  loss_logits_1: 0.5541 (0.5886)  loss_logits_2: 0.5363 (0.5809)  loss_logits_3: 0.5369 (0.5666)  loss_logits_4: 0.5127 (0.5601)  loss_logits_dn_0: 0.3044 (0.3039)  loss_logits_dn_1: 0.2849 (0.2905)  loss_logits_dn_2: 0.2869 (0.2873)  loss_logits_dn_3: 0.2857 (0.2866)  loss_logits_dn_4: 0.2924 (0.2872)  loss_logits_dn_5: 0.2970 (0.2885)  loss_logits_interm: 0.7093 (0.6627)
Test:  [ 0/25]  eta: 0:00:44  loss: 0.6208 (0.6208)  loss_line: 0.1185 (0.1185)  loss_logits: 0.5023 (0.5023)  time: 1.7869  data: 1.0693  max mem: 11997
Test:  [24/25]  eta: 0:00:00  loss: 0.7385 (0.7508)  loss_line: 0.1479 (0.1493)  loss_logits: 0.5879 (0.6015)  time: 0.6370  data: 0.0407  max mem: 11997
Test: Total time: 0:00:17 (0.7006 s / it)
Averaged stats: loss: 0.7385 (0.7508)  loss_line: 0.1479 (0.1493)  loss_logits: 0.5879 (0.6015)
Epoch: [8]  [  0/422]  eta: 0:18:25  lr: 0.000013  loss: 31.9787 (31.9787)  loss_line: 0.1207 (0.1207)  loss_line_0: 0.1765 (0.1765)  loss_line_1: 0.1608 (0.1608)  loss_line_2: 0.1288 (0.1288)  loss_line_3: 0.1349 (0.1349)  loss_line_4: 0.1224 (0.1224)  loss_line_dn_0: 4.1384 (4.1384)  loss_line_dn_1: 4.1245 (4.1245)  loss_line_dn_2: 4.1197 (4.1197)  loss_line_dn_3: 4.1183 (4.1183)  loss_line_dn_4: 4.1130 (4.1130)  loss_line_dn_5: 4.1122 (4.1122)  loss_line_interm: 0.2000 (0.2000)  loss_logits: 0.7090 (0.7090)  loss_logits_0: 0.5529 (0.5529)  loss_logits_1: 0.5650 (0.5650)  loss_logits_2: 0.6576 (0.6576)  loss_logits_3: 0.6317 (0.6317)  loss_logits_4: 0.6873 (0.6873)  loss_logits_dn_0: 0.2791 (0.2791)  loss_logits_dn_1: 0.2852 (0.2852)  loss_logits_dn_2: 0.2854 (0.2854)  loss_logits_dn_3: 0.2828 (0.2828)  loss_logits_dn_4: 0.2888 (0.2888)  loss_logits_dn_5: 0.2966 (0.2966)  loss_logits_interm: 0.6872 (0.6872)  time: 2.6208  data: 1.5029  max mem: 11997
Epoch: [8]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 40.1152 (37.4190)  loss_line: 0.1735 (0.3443)  loss_line_0: 0.2050 (0.3722)  loss_line_1: 0.2210 (0.3691)  loss_line_2: 0.1901 (0.3561)  loss_line_3: 0.1780 (0.3473)  loss_line_4: 0.1713 (0.3440)  loss_line_dn_0: 5.2045 (4.8368)  loss_line_dn_1: 5.2005 (4.8243)  loss_line_dn_2: 5.2026 (4.8225)  loss_line_dn_3: 5.2025 (4.8219)  loss_line_dn_4: 5.2022 (4.8218)  loss_line_dn_5: 5.2021 (4.8219)  loss_line_interm: 0.2506 (0.4166)  loss_logits: 0.5370 (0.5643)  loss_logits_0: 0.5774 (0.6062)  loss_logits_1: 0.5812 (0.6048)  loss_logits_2: 0.5701 (0.5933)  loss_logits_3: 0.5388 (0.5843)  loss_logits_4: 0.5228 (0.5740)  loss_logits_dn_0: 0.3127 (0.3018)  loss_logits_dn_1: 0.3050 (0.2873)  loss_logits_dn_2: 0.3040 (0.2830)  loss_logits_dn_3: 0.3000 (0.2820)  loss_logits_dn_4: 0.3000 (0.2841)  loss_logits_dn_5: 0.3020 (0.2857)  loss_logits_interm: 0.6624 (0.6696)  time: 1.1896  data: 0.0085  max mem: 11997
Epoch: [8] Total time: 0:08:15 (1.1738 s / it)
Averaged stats: lr: 0.000013  loss: 40.1152 (37.4190)  loss_line: 0.1735 (0.3443)  loss_line_0: 0.2050 (0.3722)  loss_line_1: 0.2210 (0.3691)  loss_line_2: 0.1901 (0.3561)  loss_line_3: 0.1780 (0.3473)  loss_line_4: 0.1713 (0.3440)  loss_line_dn_0: 5.2045 (4.8368)  loss_line_dn_1: 5.2005 (4.8243)  loss_line_dn_2: 5.2026 (4.8225)  loss_line_dn_3: 5.2025 (4.8219)  loss_line_dn_4: 5.2022 (4.8218)  loss_line_dn_5: 5.2021 (4.8219)  loss_line_interm: 0.2506 (0.4166)  loss_logits: 0.5370 (0.5643)  loss_logits_0: 0.5774 (0.6062)  loss_logits_1: 0.5812 (0.6048)  loss_logits_2: 0.5701 (0.5933)  loss_logits_3: 0.5388 (0.5843)  loss_logits_4: 0.5228 (0.5740)  loss_logits_dn_0: 0.3127 (0.3018)  loss_logits_dn_1: 0.3050 (0.2873)  loss_logits_dn_2: 0.3040 (0.2830)  loss_logits_dn_3: 0.3000 (0.2820)  loss_logits_dn_4: 0.3000 (0.2841)  loss_logits_dn_5: 0.3020 (0.2857)  loss_logits_interm: 0.6624 (0.6696)
Test:  [ 0/25]  eta: 0:00:43  loss: 0.5980 (0.5980)  loss_line: 0.1543 (0.1543)  loss_logits: 0.4437 (0.4437)  time: 1.7470  data: 1.0264  max mem: 11997
Test:  [24/25]  eta: 0:00:00  loss: 0.6829 (0.7093)  loss_line: 0.1458 (0.1551)  loss_logits: 0.5286 (0.5542)  time: 0.6406  data: 0.0413  max mem: 11997
Test: Total time: 0:00:17 (0.6950 s / it)
Averaged stats: loss: 0.6829 (0.7093)  loss_line: 0.1458 (0.1551)  loss_logits: 0.5286 (0.5542)
Epoch: [9]  [  0/422]  eta: 0:19:39  lr: 0.000013  loss: 39.6805 (39.6805)  loss_line: 0.2487 (0.2487)  loss_line_0: 0.2862 (0.2862)  loss_line_1: 0.2985 (0.2985)  loss_line_2: 0.2665 (0.2665)  loss_line_3: 0.2605 (0.2605)  loss_line_4: 0.2519 (0.2519)  loss_line_dn_0: 5.3650 (5.3650)  loss_line_dn_1: 5.3571 (5.3571)  loss_line_dn_2: 5.3584 (5.3584)  loss_line_dn_3: 5.3573 (5.3573)  loss_line_dn_4: 5.3567 (5.3567)  loss_line_dn_5: 5.3563 (5.3563)  loss_line_interm: 0.3109 (0.3109)  loss_logits: 0.5422 (0.5422)  loss_logits_0: 0.5389 (0.5389)  loss_logits_1: 0.5003 (0.5003)  loss_logits_2: 0.5136 (0.5136)  loss_logits_3: 0.5333 (0.5333)  loss_logits_4: 0.5537 (0.5537)  loss_logits_dn_0: 0.3224 (0.3224)  loss_logits_dn_1: 0.3165 (0.3165)  loss_logits_dn_2: 0.3172 (0.3172)  loss_logits_dn_3: 0.3080 (0.3080)  loss_logits_dn_4: 0.3071 (0.3071)  loss_logits_dn_5: 0.3107 (0.3107)  loss_logits_interm: 0.5426 (0.5426)  time: 2.7943  data: 1.5767  max mem: 11997
Epoch: [9]  [421/422]  eta: 0:00:01  lr: 0.000013  loss: 38.0286 (37.8497)  loss_line: 0.1524 (0.3398)  loss_line_0: 0.1739 (0.3656)  loss_line_1: 0.1738 (0.3633)  loss_line_2: 0.1632 (0.3520)  loss_line_3: 0.1550 (0.3442)  loss_line_4: 0.1528 (0.3406)  loss_line_dn_0: 5.2153 (4.9597)  loss_line_dn_1: 5.2015 (4.9481)  loss_line_dn_2: 5.2000 (4.9457)  loss_line_dn_3: 5.1999 (4.9450)  loss_line_dn_4: 5.1995 (4.9445)  loss_line_dn_5: 5.1990 (4.9446)  loss_line_interm: 0.2306 (0.3990)  loss_logits: 0.5154 (0.5347)  loss_logits_0: 0.5189 (0.5828)  loss_logits_1: 0.5171 (0.5690)  loss_logits_2: 0.5269 (0.5566)  loss_logits_3: 0.5108 (0.5470)  loss_logits_4: 0.5217 (0.5377)  loss_logits_dn_0: 0.3101 (0.2986)  loss_logits_dn_1: 0.3054 (0.2816)  loss_logits_dn_2: 0.3010 (0.2761)  loss_logits_dn_3: 0.3010 (0.2761)  loss_logits_dn_4: 0.3033 (0.2770)  loss_logits_dn_5: 0.3009 (0.2777)  loss_logits_interm: 0.5940 (0.6429)  time: 1.1614  data: 0.0092  max mem: 11997
Epoch: [9] Total time: 0:08:22 (1.1913 s / it)
Averaged stats: lr: 0.000013  loss: 38.0286 (37.8497)  loss_line: 0.1524 (0.3398)  loss_line_0: 0.1739 (0.3656)  loss_line_1: 0.1738 (0.3633)  loss_line_2: 0.1632 (0.3520)  loss_line_3: 0.1550 (0.3442)  loss_line_4: 0.1528 (0.3406)  loss_line_dn_0: 5.2153 (4.9597)  loss_line_dn_1: 5.2015 (4.9481)  loss_line_dn_2: 5.2000 (4.9457)  loss_line_dn_3: 5.1999 (4.9450)  loss_line_dn_4: 5.1995 (4.9445)  loss_line_dn_5: 5.1990 (4.9446)  loss_line_interm: 0.2306 (0.3990)  loss_logits: 0.5154 (0.5347)  loss_logits_0: 0.5189 (0.5828)  loss_logits_1: 0.5171 (0.5690)  loss_logits_2: 0.5269 (0.5566)  loss_logits_3: 0.5108 (0.5470)  loss_logits_4: 0.5217 (0.5377)  loss_logits_dn_0: 0.3101 (0.2986)  loss_logits_dn_1: 0.3054 (0.2816)  loss_logits_dn_2: 0.3010 (0.2761)  loss_logits_dn_3: 0.3010 (0.2761)  loss_logits_dn_4: 0.3033 (0.2770)  loss_logits_dn_5: 0.3009 (0.2777)  loss_logits_interm: 0.5940 (0.6429)
Test:  [ 0/25]  eta: 0:00:49  loss: 0.5266 (0.5266)  loss_line: 0.1051 (0.1051)  loss_logits: 0.4215 (0.4215)  time: 1.9840  data: 1.2346  max mem: 11997
Test:  [24/25]  eta: 0:00:00  loss: 0.6228 (0.6462)  loss_line: 0.1415 (0.1404)  loss_logits: 0.4845 (0.5058)  time: 0.6396  data: 0.0436  max mem: 11997
Test: Total time: 0:00:17 (0.7057 s / it)
Averaged stats: loss: 0.6228 (0.6462)  loss_line: 0.1415 (0.1404)  loss_logits: 0.4845 (0.5058)
Epoch: [10]  [  0/422]  eta: 0:19:47  lr: 0.000013  loss: 28.8314 (28.8314)  loss_line: 0.1406 (0.1406)  loss_line_0: 0.1543 (0.1543)  loss_line_1: 0.1507 (0.1507)  loss_line_2: 0.1186 (0.1186)  loss_line_3: 0.1239 (0.1239)  loss_line_4: 0.1262 (0.1262)  loss_line_dn_0: 3.7109 (3.7109)  loss_line_dn_1: 3.7019 (3.7019)  loss_line_dn_2: 3.7000 (3.7000)  loss_line_dn_3: 3.6991 (3.6991)  loss_line_dn_4: 3.6980 (3.6980)  loss_line_dn_5: 3.6975 (3.6975)  loss_line_interm: 0.2150 (0.2150)  loss_logits: 0.4811 (0.4811)  loss_logits_0: 0.5385 (0.5385)  loss_logits_1: 0.5314 (0.5314)  loss_logits_2: 0.5019 (0.5019)  loss_logits_3: 0.5182 (0.5182)  loss_logits_4: 0.4939 (0.4939)  loss_logits_dn_0: 0.3257 (0.3257)  loss_logits_dn_1: 0.3221 (0.3221)  loss_logits_dn_2: 0.3182 (0.3182)  loss_logits_dn_3: 0.3193 (0.3193)  loss_logits_dn_4: 0.3202 (0.3202)  loss_logits_dn_5: 0.3204 (0.3204)  loss_logits_interm: 0.6038 (0.6038)  time: 2.8128  data: 1.1073  max mem: 11997
Loss is nan, stopping training
{'loss_line': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_0': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_1': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_2': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_3': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_4': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_dn_0': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_dn_1': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_dn_2': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_dn_3': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_dn_4': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_dn_5': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_line_interm': tensor(0., device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits': tensor(0.0106, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_0': tensor(0.0461, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_1': tensor(0.0132, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_2': tensor(0.0100, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_3': tensor(0.0191, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_4': tensor(0.0137, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_dn_0': tensor(nan, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_dn_1': tensor(nan, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_dn_2': tensor(nan, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_dn_3': tensor(nan, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_dn_4': tensor(nan, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_dn_5': tensor(nan, device='cuda:0', grad_fn=<MulBackward0>), 'loss_logits_interm': tensor(0.0247, device='cuda:0', grad_fn=<MulBackward0>)}
[rank0]:[W907 07:04:33.437642172 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
E0907 07:04:35.578000 10051 torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 0 (pid: 10062) of binary: /usr/bin/python3
Traceback (most recent call last):
  File "/usr/local/bin/torchrun", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 357, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 901, in main
    run(args)
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/run.py", line 892, in run
    elastic_launch(
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 143, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/distributed/launcher/api.py", line 277, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
main.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-09-07_07:04:35
  host      : 5eadc0872773
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 10062)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================ ```

[SebastianJanampa]

Hi @Ibrah-N ,
I'm sorry for not getting back to you sooner. I haven't encountered that problem in my experiments, but after reviewing your data logs, I believe the issue is related to the training batch size. In my experiments, I used a batch size of 16. If you use a smaller batch size, you should also use a smaller learning rate (you can use either linear or squared scaling rules to adjust the learning rate).

I can also observe that the logits for the denoising line segments are the ones producing nan values. A way to address this problem is by using float32 for the classification loss.

Also, are you training your model from scratch, or are you using the pretrained weights from this repo?