From 1ca8b0b9e0613c1e02f6a5d8b49e20c4d6916412 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Thu, 22 Apr 2021 08:15:58 +0200 Subject: Fixed training script, able to train vqvae --- text_recognizer/networks/vqvae/vqvae.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'text_recognizer/networks/vqvae/vqvae.py') diff --git a/text_recognizer/networks/vqvae/vqvae.py b/text_recognizer/networks/vqvae/vqvae.py index 50448b4..1f08e5e 100644 --- a/text_recognizer/networks/vqvae/vqvae.py +++ b/text_recognizer/networks/vqvae/vqvae.py @@ -1,8 +1,7 @@ """The VQ-VAE.""" -from typing import List, Optional, Tuple, Type +from typing import Any, Dict, List, Optional, Tuple -import torch from torch import nn from torch import Tensor @@ -25,6 +24,8 @@ class VQVAE(nn.Module): beta: float = 0.25, activation: str = "leaky_relu", dropout_rate: float = 0.0, + *args: Any, + **kwargs: Dict, ) -> None: super().__init__() -- cgit v1.2.3-70-g09d2