[2024-03-07 13:29:59,877] INFO: Will use torch.nn.parallel.DistributedDataParallel() and 4 gpus
[2024-03-07 13:29:59,894] INFO: NVIDIA A10
[2024-03-07 13:29:59,894] INFO: NVIDIA A10
[2024-03-07 13:29:59,894] INFO: NVIDIA A10
[2024-03-07 13:29:59,894] INFO: NVIDIA A10
[2024-03-07 13:30:07,464] INFO: using attention_type=efficient
[2024-03-07 13:30:07,467] INFO: using attention_type=efficient
[2024-03-07 13:30:07,471] INFO: using attention_type=efficient
[2024-03-07 13:30:07,474] INFO: using attention_type=efficient
[2024-03-07 13:30:07,477] INFO: using attention_type=efficient
[2024-03-07 13:30:07,481] INFO: using attention_type=efficient
[2024-03-07 13:30:11,211] INFO: DistributedDataParallel(
  (module): MLPF(
    (nn0): Sequential(
      (0): Linear(in_features=42, out_features=256, bias=True)
      (1): ELU(alpha=1.0)
      (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
      (3): Dropout(p=0.3, inplace=False)
      (4): Linear(in_features=256, out_features=256, bias=True)
    )
    (conv_id): ModuleList(
      (0-2): 3 x SelfAttentionLayer(
        (mha): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
        )
        (norm0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (seq): Sequential(
          (0): Linear(in_features=256, out_features=256, bias=True)
          (1): ELU(alpha=1.0)
          (2): Linear(in_features=256, out_features=256, bias=True)
          (3): ELU(alpha=1.0)
        )
        (dropout): Dropout(p=0.3, inplace=False)
      )
    )
    (conv_reg): ModuleList(
      (0-2): 3 x SelfAttentionLayer(
        (mha): MultiheadAttention(
          (out_proj): NonDynamicallyQuantizableLinear(in_features=256, out_features=256, bias=True)
        )
        (norm0): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (norm1): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (seq): Sequential(
          (0): Linear(in_features=256, out_features=256, bias=True)
          (1): ELU(alpha=1.0)
          (2): Linear(in_features=256, out_features=256, bias=True)
          (3): ELU(alpha=1.0)
        )
        (dropout): Dropout(p=0.3, inplace=False)
      )
    )
    (nn_id): Sequential(
      (0): Linear(in_features=810, out_features=256, bias=True)
      (1): ELU(alpha=1.0)
      (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
      (3): Dropout(p=0.3, inplace=False)
      (4): Linear(in_features=256, out_features=9, bias=True)
    )
    (nn_pt): RegressionOutput(
      (nn): Sequential(
        (0): Linear(in_features=819, out_features=256, bias=True)
        (1): ELU(alpha=1.0)
        (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (3): Dropout(p=0.3, inplace=False)
        (4): Linear(in_features=256, out_features=2, bias=True)
      )
    )
    (nn_eta): RegressionOutput(
      (nn): Sequential(
        (0): Linear(in_features=819, out_features=256, bias=True)
        (1): ELU(alpha=1.0)
        (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (3): Dropout(p=0.3, inplace=False)
        (4): Linear(in_features=256, out_features=2, bias=True)
      )
    )
    (nn_sin_phi): RegressionOutput(
      (nn): Sequential(
        (0): Linear(in_features=819, out_features=256, bias=True)
        (1): ELU(alpha=1.0)
        (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (3): Dropout(p=0.3, inplace=False)
        (4): Linear(in_features=256, out_features=2, bias=True)
      )
    )
    (nn_cos_phi): RegressionOutput(
      (nn): Sequential(
        (0): Linear(in_features=819, out_features=256, bias=True)
        (1): ELU(alpha=1.0)
        (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (3): Dropout(p=0.3, inplace=False)
        (4): Linear(in_features=256, out_features=2, bias=True)
      )
    )
    (nn_energy): RegressionOutput(
      (nn): Sequential(
        (0): Linear(in_features=819, out_features=256, bias=True)
        (1): ELU(alpha=1.0)
        (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
        (3): Dropout(p=0.3, inplace=False)
        (4): Linear(in_features=256, out_features=2, bias=True)
      )
    )
    (nn_charge): Sequential(
      (0): Linear(in_features=819, out_features=256, bias=True)
      (1): ELU(alpha=1.0)
      (2): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
      (3): Dropout(p=0.3, inplace=False)
      (4): Linear(in_features=256, out_features=3, bias=True)
    )
  )
)
[2024-03-07 13:30:11,212] INFO: Trainable parameters: 3928342
[2024-03-07 13:30:11,212] INFO: Non-trainable parameters: 0
[2024-03-07 13:30:11,212] INFO: Total parameters: 3928342
[2024-03-07 13:30:11,218] INFO:                               Modules Trainable params Non-tranable params  Trainable Parameters Non-tranable Parameters
                  module.nn0.0.weight              NaN                 NaN               10752.0                       -
                    module.nn0.0.bias              NaN                 NaN                 256.0                       -
                  module.nn0.2.weight              NaN                 NaN                 256.0                       -
                    module.nn0.2.bias              NaN                 NaN                 256.0                       -
                  module.nn0.4.weight              NaN                 NaN               65536.0                       -
                    module.nn0.4.bias              NaN                 NaN                 256.0                       -
  module.conv_id.0.mha.in_proj_weight              NaN                 NaN              196608.0                       -
    module.conv_id.0.mha.in_proj_bias              NaN                 NaN                 768.0                       -
 module.conv_id.0.mha.out_proj.weight              NaN                 NaN               65536.0                       -
   module.conv_id.0.mha.out_proj.bias              NaN                 NaN                 256.0                       -
        module.conv_id.0.norm0.weight              NaN                 NaN                 256.0                       -
          module.conv_id.0.norm0.bias              NaN                 NaN                 256.0                       -
        module.conv_id.0.norm1.weight              NaN                 NaN                 256.0                       -
          module.conv_id.0.norm1.bias              NaN                 NaN                 256.0                       -
        module.conv_id.0.seq.0.weight              NaN                 NaN               65536.0                       -
          module.conv_id.0.seq.0.bias              NaN                 NaN                 256.0                       -
        module.conv_id.0.seq.2.weight              NaN                 NaN               65536.0                       -
          module.conv_id.0.seq.2.bias              NaN                 NaN                 256.0                       -
  module.conv_id.1.mha.in_proj_weight              NaN                 NaN              196608.0                       -
    module.conv_id.1.mha.in_proj_bias              NaN                 NaN                 768.0                       -
 module.conv_id.1.mha.out_proj.weight              NaN                 NaN               65536.0                       -
   module.conv_id.1.mha.out_proj.bias              NaN                 NaN                 256.0                       -
        module.conv_id.1.norm0.weight              NaN                 NaN                 256.0                       -
          module.conv_id.1.norm0.bias              NaN                 NaN                 256.0                       -
        module.conv_id.1.norm1.weight              NaN                 NaN                 256.0                       -
          module.conv_id.1.norm1.bias              NaN                 NaN                 256.0                       -
        module.conv_id.1.seq.0.weight              NaN                 NaN               65536.0                       -
          module.conv_id.1.seq.0.bias              NaN                 NaN                 256.0                       -
        module.conv_id.1.seq.2.weight              NaN                 NaN               65536.0                       -
          module.conv_id.1.seq.2.bias              NaN                 NaN                 256.0                       -
  module.conv_id.2.mha.in_proj_weight              NaN                 NaN              196608.0                       -
    module.conv_id.2.mha.in_proj_bias              NaN                 NaN                 768.0                       -
 module.conv_id.2.mha.out_proj.weight              NaN                 NaN               65536.0                       -
   module.conv_id.2.mha.out_proj.bias              NaN                 NaN                 256.0                       -
        module.conv_id.2.norm0.weight              NaN                 NaN                 256.0                       -
          module.conv_id.2.norm0.bias              NaN                 NaN                 256.0                       -
        module.conv_id.2.norm1.weight              NaN                 NaN                 256.0                       -
          module.conv_id.2.norm1.bias              NaN                 NaN                 256.0                       -
        module.conv_id.2.seq.0.weight              NaN                 NaN               65536.0                       -
          module.conv_id.2.seq.0.bias              NaN                 NaN                 256.0                       -
        module.conv_id.2.seq.2.weight              NaN                 NaN               65536.0                       -
          module.conv_id.2.seq.2.bias              NaN                 NaN                 256.0                       -
 module.conv_reg.0.mha.in_proj_weight              NaN                 NaN              196608.0                       -
   module.conv_reg.0.mha.in_proj_bias              NaN                 NaN                 768.0                       -
module.conv_reg.0.mha.out_proj.weight              NaN                 NaN               65536.0                       -
  module.conv_reg.0.mha.out_proj.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.0.norm0.weight              NaN                 NaN                 256.0                       -
         module.conv_reg.0.norm0.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.0.norm1.weight              NaN                 NaN                 256.0                       -
         module.conv_reg.0.norm1.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.0.seq.0.weight              NaN                 NaN               65536.0                       -
         module.conv_reg.0.seq.0.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.0.seq.2.weight              NaN                 NaN               65536.0                       -
         module.conv_reg.0.seq.2.bias              NaN                 NaN                 256.0                       -
 module.conv_reg.1.mha.in_proj_weight              NaN                 NaN              196608.0                       -
   module.conv_reg.1.mha.in_proj_bias              NaN                 NaN                 768.0                       -
module.conv_reg.1.mha.out_proj.weight              NaN                 NaN               65536.0                       -
  module.conv_reg.1.mha.out_proj.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.1.norm0.weight              NaN                 NaN                 256.0                       -
         module.conv_reg.1.norm0.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.1.norm1.weight              NaN                 NaN                 256.0                       -
         module.conv_reg.1.norm1.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.1.seq.0.weight              NaN                 NaN               65536.0                       -
         module.conv_reg.1.seq.0.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.1.seq.2.weight              NaN                 NaN               65536.0                       -
         module.conv_reg.1.seq.2.bias              NaN                 NaN                 256.0                       -
 module.conv_reg.2.mha.in_proj_weight              NaN                 NaN              196608.0                       -
   module.conv_reg.2.mha.in_proj_bias              NaN                 NaN                 768.0                       -
module.conv_reg.2.mha.out_proj.weight              NaN                 NaN               65536.0                       -
  module.conv_reg.2.mha.out_proj.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.2.norm0.weight              NaN                 NaN                 256.0                       -
         module.conv_reg.2.norm0.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.2.norm1.weight              NaN                 NaN                 256.0                       -
         module.conv_reg.2.norm1.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.2.seq.0.weight              NaN                 NaN               65536.0                       -
         module.conv_reg.2.seq.0.bias              NaN                 NaN                 256.0                       -
       module.conv_reg.2.seq.2.weight              NaN                 NaN               65536.0                       -
         module.conv_reg.2.seq.2.bias              NaN                 NaN                 256.0                       -
                module.nn_id.0.weight              NaN                 NaN              207360.0                       -
                  module.nn_id.0.bias              NaN                 NaN                 256.0                       -
                module.nn_id.2.weight              NaN                 NaN                 256.0                       -
                  module.nn_id.2.bias              NaN                 NaN                 256.0                       -
                module.nn_id.4.weight              NaN                 NaN                2304.0                       -
                  module.nn_id.4.bias              NaN                 NaN                   9.0                       -
             module.nn_pt.nn.0.weight              NaN                 NaN              209664.0                       -
               module.nn_pt.nn.0.bias              NaN                 NaN                 256.0                       -
             module.nn_pt.nn.2.weight              NaN                 NaN                 256.0                       -
               module.nn_pt.nn.2.bias              NaN                 NaN                 256.0                       -
             module.nn_pt.nn.4.weight              NaN                 NaN                 512.0                       -
               module.nn_pt.nn.4.bias              NaN                 NaN                   2.0                       -
            module.nn_eta.nn.0.weight              NaN                 NaN              209664.0                       -
              module.nn_eta.nn.0.bias              NaN                 NaN                 256.0                       -
            module.nn_eta.nn.2.weight              NaN                 NaN                 256.0                       -
              module.nn_eta.nn.2.bias              NaN                 NaN                 256.0                       -
            module.nn_eta.nn.4.weight              NaN                 NaN                 512.0                       -
              module.nn_eta.nn.4.bias              NaN                 NaN                   2.0                       -
        module.nn_sin_phi.nn.0.weight              NaN                 NaN              209664.0                       -
          module.nn_sin_phi.nn.0.bias              NaN                 NaN                 256.0                       -
        module.nn_sin_phi.nn.2.weight              NaN                 NaN                 256.0                       -
          module.nn_sin_phi.nn.2.bias              NaN                 NaN                 256.0                       -
        module.nn_sin_phi.nn.4.weight              NaN                 NaN                 512.0                       -
          module.nn_sin_phi.nn.4.bias              NaN                 NaN                   2.0                       -
        module.nn_cos_phi.nn.0.weight              NaN                 NaN              209664.0                       -
          module.nn_cos_phi.nn.0.bias              NaN                 NaN                 256.0                       -
        module.nn_cos_phi.nn.2.weight              NaN                 NaN                 256.0                       -
          module.nn_cos_phi.nn.2.bias              NaN                 NaN                 256.0                       -
        module.nn_cos_phi.nn.4.weight              NaN                 NaN                 512.0                       -
          module.nn_cos_phi.nn.4.bias              NaN                 NaN                   2.0                       -
         module.nn_energy.nn.0.weight              NaN                 NaN              209664.0                       -
           module.nn_energy.nn.0.bias              NaN                 NaN                 256.0                       -
         module.nn_energy.nn.2.weight              NaN                 NaN                 256.0                       -
           module.nn_energy.nn.2.bias              NaN                 NaN                 256.0                       -
         module.nn_energy.nn.4.weight              NaN                 NaN                 512.0                       -
           module.nn_energy.nn.4.bias              NaN                 NaN                   2.0                       -
            module.nn_charge.0.weight              NaN                 NaN              209664.0                       -
              module.nn_charge.0.bias              NaN                 NaN                 256.0                       -
            module.nn_charge.2.weight              NaN                 NaN                 256.0                       -
              module.nn_charge.2.bias              NaN                 NaN                 256.0                       -
            module.nn_charge.4.weight              NaN                 NaN                 768.0                       -
              module.nn_charge.4.bias              NaN                 NaN                   3.0                       -
[2024-03-07 13:30:11,306] INFO: Creating experiment dir /pfvol/experiments/MLPF_cms_Transformer_MET_Falsepyg-cms-small_20240307_132958_158969
[2024-03-07 13:30:11,306] INFO: Model directory /pfvol/experiments/MLPF_cms_Transformer_MET_Falsepyg-cms-small_20240307_132958_158969