From 0b0ecc94f9bf3cd8c582ac06c6b3f145f130a0e4 Mon Sep 17 00:00:00 2001 From: Kashif Rasul Date: Mon, 6 Jun 2022 11:45:31 +0200 Subject: [PATCH] fix typo --- switch/estimator.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/switch/estimator.py b/switch/estimator.py index cab46f6..0954c92 100644 --- a/switch/estimator.py +++ b/switch/estimator.py @@ -108,8 +108,7 @@ class SwitchTransformerEstimator(PyTorchLightningEstimator): self.capacity_factor = capacity_factor self.is_scale_prob = is_scale_prob self.drop_tokens = drop_tokens - - + self.num_feat_dynamic_real = num_feat_dynamic_real self.num_feat_static_cat = num_feat_static_cat self.num_feat_static_real = num_feat_static_real @@ -321,4 +320,4 @@ class SwitchTransformerEstimator(PyTorchLightningEstimator): num_parallel_samples=self.num_parallel_samples, ) - return TransformerLightningModule(model=model, loss=self.loss) + return SwitchTransformerLightningModule(model=model, loss=self.loss)