From f7ef47738e4cb4a07e5f734d50ad306aec32e7ff Mon Sep 17 00:00:00 2001
From: David Rotermund <54365609+davrot@users.noreply.github.com>
Date: Tue, 25 Jul 2023 00:57:03 +0200
Subject: [PATCH] Add files via upload

---
 functions/SoftmaxPower.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/functions/SoftmaxPower.py b/functions/SoftmaxPower.py
index 8a39ddc..b0fd0ce 100644
--- a/functions/SoftmaxPower.py
+++ b/functions/SoftmaxPower.py
@@ -32,7 +32,11 @@ class SoftmaxPower(torch.nn.Module):
             self.no_input_mode = False
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
-        output: torch.Tensor = torch.abs(input).pow(self.power)
+        if self.power != 0.0:
+            output: torch.Tensor = torch.abs(input).pow(self.power)
+        else:
+            output: torch.Tensor = torch.exp(input)
+
         if self.dim is None:
             output = output / output.sum()
         else: