gorarakelyan

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from datasets import load_dataset

from aim import Run
from aim.hf_dataset import HFDataset
# from aim.pytorch import track_gradients_dists, track_params_dists

gorarakelyan

import random
import time

from aim import Run
from aim import TrainingFlow
from aim.sdk.callbacks import events


# Define the callbacks
class MyCallbacks:

gorarakelyan

import click
from datetime import timedelta

from aim import Repo


@click.command()
@click.option('--delta', prompt='Time delta in hours from your current timezone to UTC', required=True, type=int)
@click.option('--repo', required=False, type=click.Path(exists=True,
                                                        file_okay=False,

gorarakelyan

"""MNIST example.
Library file which executes the training and evaluation loop for MNIST.
The data is loaded using tensorflow_datasets.
"""

# See issue #620.
# pytype: disable=wrong-keyword-args

from absl import logging
from flax import linen as nn