Codebase Refactor

Author: federicobrancasi

DeepQuant/CustomForwards/Activations.py

@@ -5,8 +5,8 @@
 # Federico Brancasi <[email protected]>
 
 import torch.nn as nn
-from torch import Tensor
 from brevitas.nn.quant_layer import QuantNonLinearActLayer
+from torch import Tensor
 
 
 class WrapperActivation(nn.Module):
@@ -28,4 +28,4 @@ def activationForward(self: QuantNonLinearActLayer, inp: Tensor) -> Tensor:
     else:
         output = quantInput
     quantOutput = self.act_quant(output) if self.act_quant is not None else output
-    return quantOutput
+    return quantOutput

DeepQuant/CustomForwards/Linear.py

@@ -5,8 +5,8 @@
 # Federico Brancasi <[email protected]>
 
 import torch.nn as nn
-from torch import Tensor
 from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer
+from torch import Tensor
 
 
 class WrapperLinear(nn.Module):
@@ -33,4 +33,4 @@ def linearForward(self: QuantWeightBiasInputOutputLayer, inp: Tensor) -> Tensor:
 
     output = self.wrappedInnerForwardImpl(quantInput, quantWeight, quantBias)
     quantOutput = self.output_quant(output)
-    return quantOutput
+    return quantOutput

DeepQuant/CustomForwards/MultiHeadAttention.py

@@ -5,10 +5,11 @@
 # Federico Brancasi <[email protected]>
 
 import math
+
 import torch
 import torch.nn.functional as F
-from torch import Tensor
 from brevitas.nn.quant_mha import QuantMultiheadAttention
+from torch import Tensor
 
 
 def mhaForward(
@@ -59,4 +60,4 @@ def mhaForward(
     )
 
     attnOutput = self.out_proj(attnOutput)
-    return attnOutput
+    return attnOutput