Update pocs/transf.py

Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com>
leonvanbokhorst · Dec 13, 2024 · 173a06e · 173a06e
1 parent 657d3e5
commit 173a06e
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/pocs/transf.py b/pocs/transf.py
@@ -291,9 +291,8 @@ def test_addition(model, num1, num2, max_digits=3):
         predicted = torch.argmax(output, dim=-1)
 
         # Filter out padding tokens
-        result = ''.join(str(x) for x in predicted[0].tolist() if x != 0).lstrip('0')
-        if not result:  # If result is empty after filtering
-            result = '0'
+        result = ''.join(str(x) for x in predicted[0].tolist() if x != 0).lstrip('0') or '0'
+
 
         print(f"{num1} + {num2} = {result} (expected {num1 + num2})")
         return result