public static VisionModel V2S() {
        return V2S("model/EfficientNet/efficientnet_v2_s.pt");
    }
@Test
    public void train() throws IOException {
        Device device = Device.CUDA((byte) 1);
        device.setDefaultDevice();

        // half precision to lower memory usage.
        var dtype = ScalarType.BFloat16;

        var model = EfficientNet.V2S();
        model.to(device, dtype);

        var transform = Transform.classification(384, 384);
        var data = new ImageDataset(64, "../imagenet/train", transform, ImageNet.folder2Target);
        var test = new ImageDataset(16, "../imagenet/val",   transform, ImageNet.folder2Target);

        var schedule = TimeFunction.piecewise(new int[] { 50000 },
                TimeFunction.linear(0.0001, 50000, 0.01),
                TimeFunction.cosine(0.0001, 50000, 0.01));
        model.setLearningRateSchedule(schedule);
        // Use parameters from the paper, the rests are Keras default values.
        // Note that Keras has different default values from PyTorch (e.g. alpha and eps).
        Optimizer optimizer = Optimizer.RMSprop(model, 0.0001, 0.9, 1E-07, 1E-05, 0.9, false);
        model.train(5, optimizer, Loss.nll(), data, test, null, new Accuracy());
    }