code: import os
  import torch
  from weaver.utils.logger import _logger
  from weaver.utils.import_tools import import_module
  
  ParticleTransformer = import_module(
      os.path.join(os.path.dirname(__file__), 'ParticleTransformer.py'), 'ParT').ParticleTransformer
  
  
  class ParticleTransformerWrapper(torch.nn.Module):
      def __init__(self, **kwargs) -> None:
          super().__init__()
          self.mod = ParticleTransformer(**kwargs)
  
      @torch.jit.ignore
      def no_weight_decay(self):
          return {'mod.cls_token', }
  
      def forward(self, points, features, lorentz_vectors, mask):
          return self.mod(features, v=lorentz_vectors, mask=mask)
  
  
  def get_model(data_config, **kwargs):
  
      cfg = dict(
          input_dim=len(data_config.input_dicts['pf_features']),
          num_classes=len(data_config.label_value),
          # network configurations
          pair_input_dim=4,
          embed_dims=[128, 512, 128],
          pair_embed_dims=[64, 64, 64],
          num_heads=8,
          num_layers=8,
          num_cls_layers=2,
          block_params=None,
          cls_block_params={'dropout': 0, 'attn_dropout': 0, 'activation_dropout': 0},
          fc_params=[],
          activation='gelu',
          # misc
          trim=True,
          for_inference=False,
      )
      cfg.update(**kwargs)
      _logger.info('Model config: %s' % str(cfg))
  
      model = ParticleTransformerWrapper(**cfg)
  
      model_info = {
          'input_names': list(data_config.input_names),
          'input_shapes': {k: ((1,) + s[1:]) for k, s in data_config.input_shapes.items()},
          'output_names': ['softmax'],
          'dynamic_axes': {**{k: {0: 'N', 2: 'n_' + k.split('_')[0]} for k in data_config.input_names}, **{'softmax': {0: 'N'}}},
      }
  
      return model, model_info
  
  
  def get_loss(data_config, **kwargs):
      return torch.nn.CrossEntropyLoss()
  
apis: torch.nn.CrossEntropyLoss