fix: 0 is not none

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

mha.py

@@ -646,7 +646,7 @@ class MHA(nn.Module):
         if not self.cross_attn and self.num_heads_kv == self.num_heads:
             assert x_kv is None and mixer_subset is None
             lora_kwargs = {}
-            if task:
+            if task is not None:
                 lora_kwargs['task'] = task
                 lora_kwargs['residual'] = self.return_residual
 

mlp.py

@@ -49,7 +49,7 @@ class Mlp(nn.Module):
 
     def forward(self, x, task):
         lora_kwargs = {}
-        if task:
+        if task is not None:
             lora_kwargs['task'] = task
         y = self.fc1(x, **lora_kwargs)
         y = self.activation(y)

modeling_lora.py

@@ -181,7 +181,7 @@ class LoRAParametrization(nn.Module):
 
             def new_forward(self, input, task, residual=False):
                 task_idx = adaptation_map[task] if task else None
-                if task_idx:
+                if task_idx is not None:
                     weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
                 else:
                     weights = self.weight
@@ -210,7 +210,7 @@ class LoRAParametrization(nn.Module):
 
             def new_forward(self, input, task):
                 task_idx = adaptation_map[task] if task else None
-                if task_idx:
+                if task_idx is not None:
                     weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
                 else:
                     weights = self.weight

modeling_xlm_roberta.py

@@ -314,7 +314,7 @@ class XLMRobertaPooler(nn.Module):
         # We "pool" the model by simply taking the hidden state corresponding
         # to the first token.
         lora_kwargs = {}
-        if task:
+        if task is not None:
             lora_kwargs['task'] = task
 
         first_token_tensor = hidden_states[:, 0] if pool else hidden_states
@@ -551,7 +551,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         else:
             range_iter = range(0, len(sentences), batch_size)
         lora_kwargs = {}
-        if task:
+        if task is not None:
             lora_kwargs['task'] = task
         for i in range_iter:
             encoded_input = self.tokenizer(