# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
from io import BytesIO

import numpy as np
from helpers import MLP, graph_mode

import megengine.functional as F
from megengine import load, optimizer, save
from megengine.core import TensorDict, tensor
from megengine.jit import trace
from megengine.test import assertTensorClose


def get_input():
    batch_size, input_dim = 2, 28
    data_shape, label_shape = (batch_size, input_dim), (batch_size,)
    data, label = tensor(dtype=np.float32), 
tensor(dtype=np.int32)