Bug fix for calling BatchRandTransforms

warner-benjamin · Feb 15, 2023 · 1c9168e · 1c9168e
1 parent af36a26
commit 1c9168e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/fastxtend/transform.py b/fastxtend/transform.py
@@ -56,7 +56,7 @@ def before_call(self,
     ):
         "Randomly select `self.idxs` and set `self.do` based on `self.p` if not valid `split_idx`"
         self.idxs = self.bernoulli.sample((find_bs(b),)).bool() if not split_idx and self.p<1. else torch.ones(find_bs(b)).bool()
-        self.do = self.p==1. or self.idxs.shape[-1] > 0
+        self.do = self.p==1. or self.idxs.sum() > 0
 
     def __call__(self,
         b:Tensor|tuple[Tensor,...], # Batch item(s)

diff --git a/nbs/transform.ipynb b/nbs/transform.ipynb
@@ -111,7 +111,7 @@
     "    ):\n",
     "        \"Randomly select `self.idxs` and set `self.do` based on `self.p` if not valid `split_idx`\"\n",
     "        self.idxs = self.bernoulli.sample((find_bs(b),)).bool() if not split_idx and self.p<1. else torch.ones(find_bs(b)).bool()\n",
-    "        self.do = self.p==1. or self.idxs.shape[-1] > 0\n",
+    "        self.do = self.p==1. or self.idxs.sum() > 0\n",
     "\n",
     "    def __call__(self,\n",
     "        b:Tensor|tuple[Tensor,...], # Batch item(s)\n",