common.py

#!/usr/bin/env python
# -*- encoding: utf-8 -*-

import torch

DEPTH = 2
BATCH_SIZE = 512
SEQ_LENGTH = 128
HIDDEN_SIZE = 512
NUM_CLASSES = 1000
NUM_BLOCKS = 6
IMG_SIZE = 224

def check_equal(A, B):
    return torch.allclose(A, B, rtol=1e-4, atol=1e-2)