code: def __init__(self, num_samples, num_features, **kwargs):
      super(JittorArgMaxDataset, self).__init__(**kwargs)
      self.x = jt.randn(num_samples, num_features)
      self.y = self.x.argmax(dim=-1)
      self.set_attrs(total_len=num_samples)
  
file_path: fastNLP/tests/helpers/datasets/jittor_data.py
imports: from fastNLP.envs.imports import _NEED_IMPORT_JITTOR
label_desc: Use image node_id 1 to create a new JittorArgMaxDataset object from inherited base classes: Dataset with example: obj = JittorArgMaxDataset(num_samples, num_features)