dev(narugo): fix bug on dims

_CKPT_PATTERN = re.compile(r'^monochrome-(?P<name>[a-zA-Z\d_\-]+)-(?P<epoch>\d+)\.ckpt$')

_KNOWN_MODELS = {}

_KNOWN_DIMS = {}

_KNOWN_DATASETS = {}

def _register_model(cls: Type[nn.Module], *args, name=None, **kwargs):

    name = name or cls.__model_name__

    _KNOWN_MODELS[name] = partial(cls, *args, **kwargs)

    _KNOWN_DIMS[name] = getattr(cls, '__dims__', 1)

_register_model(MonochromeAlexNet)

    )

    session_name = session_name or model_name

    model_dims = getattr(_KNOWN_MODELS[model_name], '__dims__', 1)

    model_dims = _KNOWN_DIMS[model_name]

    _log_dir = os.path.join(_LOG_DIR, session_name)

    if accelerator.is_local_main_process:

    num_workers = num_workers or min(os.cpu_count(), batch_size)

    train_dataset, test_dataset = random_split_dataset(

        full_dataset, train_size, test_size,

        trans_val=TRANSFORM2_VAL if full_dataset.__dims__ == 2 else TRANSFORM_VAL

        trans_val=TRANSFORM2_VAL if model_dims == 2 else TRANSFORM_VAL

    )

    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers,

                                  drop_last=True)