Thanks for your reply.
Is it correct to understand PyTorch eager mode as normal mode?
If so, the output of the model(input) is below. I think it is correct.
>>> out = model(dummy_input)
>>> out
(tensor([[-0.0391, -0.0840,  0.1382, -0.0397,  0.0157, -0.0448, -0.0603, -0.0823,
          0.0025, -0.0009],
        [ 0.0257,  0.0031, -0.1880,  0.0768, -0.1047, -0.0392,  0.1393, -0.0419,
         -0.0437,  0.0032],
        [ 0.0090,  0.0097, -0.0768, -0.0383, -0.0220, -0.2048, -0.1315,  0.0117,
         -0.0538, -0.0613],
        [ 0.0438, -0.1284,  0.0325,  0.0441,  0.0736,  0.1941, -0.0407, -0.0634,
          0.1074, -0.0407],
        [ 0.0440,  0.0194,  0.0147,  0.0859,  0.2149, -0.0393,  0.1640,  0.0369,
         -0.1021, -0.1820],
        [-0.3142, -0.0726, -0.0694, -0.1064, -0.1595,  0.2461,  0.1174,  0.2102,
          0.1790,  0.2188],
        [-0.0522,  0.0327, -0.1626, -0.0955,  0.0625, -0.0061,  0.0662,  0.0667,
          0.1003,  0.0635],
        [ 0.1054, -0.0456,  0.0922,  0.0559,  0.1422, -0.1924, -0.2107, -0.0572,
         -0.0424, -0.1007]], device='cuda:0', grad_fn=<AddmmBackward>), tensor([[-0.0779, -0.4483,  0.5459, -0.4263,  0.3033, -0.0147,  0.1823,  0.2561,
          0.3321, -0.8131],
        [-0.1185, -0.1932,  0.2465,  0.3930, -0.0634,  0.2440, -0.0587, -0.5931,
          0.0938,  0.2163],
        [ 0.0699, -0.2207,  0.5958, -0.0778, -0.1024, -0.1841,  0.5211,  0.0760,
          0.2308,  0.1463],
        [ 0.1858, -0.0432,  0.3188, -0.0905,  0.1415, -0.6925,  0.1487, -0.2300,
          1.0883,  0.1186],
        [-0.1471,  0.1120,  0.3354, -0.3918,  0.0748, -0.5318,  0.0106, -0.4543,
          1.2513,  0.1778],
        [ 0.4499,  0.0425,  0.3949, -0.8790, -0.1463, -0.4942, -0.4362, -0.3380,
          0.3257,  0.2104],
        [ 0.2962,  0.0098,  0.6569,  0.0520,  0.1627, -0.4044,  0.2104, -0.2278,
          0.2411,  0.0337],
        [ 0.0591, -0.0795,  0.9120, -0.5483, -0.2887, -0.2304, -0.3799, -0.5769,
          0.5903, -0.6071]], device='cuda:0', grad_fn=<AddmmBackward>))


Regarding the JIT, I don’t know much about it, so could you please tell me how to check that?
Here’s the code, is it possible to convert to ONNX even if ‘foward’ has ‘if’ in it?
class NetworkCIFAR(nn.Module):

  def __init__(self, C, num_classes, layers, auxiliary, genotype):
    super(NetworkCIFAR, self).__init__()
    self._layers = layers
    self._auxiliary = auxiliary

    stem_multiplier = 3 #RGB
    C_curr = stem_multiplier*C #C is input_channels
    self.stem = nn.Sequential(
      nn.Conv2d(3, C_curr, 3, padding=1, bias=False),
      nn.BatchNorm2d(C_curr)
    )
    
    C_prev_prev, C_prev, C_curr = C_curr, C_curr, C
    self.cells = nn.ModuleList()
    reduction_prev = False
    for i in range(layers):
      if i in [layers//3, 2*layers//3]:
        C_curr *= 2
        reduction = True
      else:
        reduction = False
      cell = Cell(genotype, C_prev_prev, C_prev, C_curr, reduction, reduction_prev)
      reduction_prev = reduction
      self.cells += [cell]
      C_prev_prev, C_prev = C_prev, cell.multiplier*C_curr
      if i == 2*layers//3:
        C_to_auxiliary = C_prev

    if auxiliary:
      self.auxiliary_head = AuxiliaryHeadCIFAR(C_to_auxiliary, num_classes)
    self.global_pooling = nn.AdaptiveAvgPool2d(1)
    self.classifier = nn.Linear(C_prev, num_classes)

  def forward(self, input):
    logits_aux = None
    s0 = s1 = self.stem(input)
    for i, cell in enumerate(self.cells):
      s0, s1 = s1, cell(s0, s1, self.drop_path_prob)
      if i == 2*self._layers//3:
        if self._auxiliary and self.training:
          logits_aux = self.auxiliary_head(s1)
    out = self.global_pooling(s1)
    logits = self.classifier(out.view(out.size(0),-1))
    return logits, logits_aux

I’m sorry for all of the questions.
Thank you for your time.