fix mask dtype comatibility

src/textbrewer/compatibility.py

@@ -0,0 +1,6 @@
+import torch
+
+if torch.__version__ < '1.2':
+    mask_dtype = torch.uint8
+else:
+    mask_dtype = torch.bool

src/textbrewer/distillation.py

@@ -13,6 +13,8 @@
 from .presets import *
 from .configurations import TrainingConfig, DistillationConfig
 
+from .compatibility import mask_dtype
+
 logger = logging.getLogger("Distillation")
 logger.setLevel(logging.INFO)
 
@@ -796,7 +798,7 @@ def _select_logits_with_mask(logits_list, masks_list):
     if len(masks_list)==len(logits_list):
         for logits,mask in zip(logits_list,masks_list):
             if len(logits.shape)==3:
-                mask = mask.unsqueeze(-1).expand_as(logits).to(torch.uint8)
+                mask = mask.unsqueeze(-1).expand_as(logits).to(mask_dtype)
                 logits_select = torch.masked_select(logits,mask).view(-1,logits.size(-1))
             else:
                 logits_select = logits #Logits_mask has no effect on logits of shape (batch_size, logits_to_be_softmaxed)
@@ -805,7 +807,7 @@ def _select_logits_with_mask(logits_list, masks_list):
         mask = masks_list[0]
         for logits in logits_list:
             if len(logits.shape)==3:
-                mask = mask.unsqueeze(-1).expand_as(logits).to(torch.uint8)
+                mask = mask.unsqueeze(-1).expand_as(logits).to(mask_dtype)
                 logits_select = torch.masked_select(logits,mask).view(-1,logits.size(-1))
             else:
                 logits_select = logits #Logits_mask has no effect on logits of shape (batch_size, logits_to_be_softmaxed)

src/textbrewer/losses.py

@@ -2,6 +2,8 @@
 import torch
 from typing import List
 
+from .compatibility import mask_dtype
+
 def kd_mse_loss(logits_S, logits_T, temperature=1):
     '''
     Calculate the mse loss between logits_S and logits_T
@@ -155,7 +157,7 @@ def cos_loss(state_S, state_T, mask=None):
         state_S = state_S.view(-1,state_S.size(-1))
         state_T = state_T.view(-1,state_T.size(-1))
     else:
-        mask = mask.to(state_S).unsqueeze(-1).expand_as(state_S).to(torch.uint8) #(bs,len,dim)
+        mask = mask.to(state_S).unsqueeze(-1).expand_as(state_S).to(mask_dtype) #(bs,len,dim)
         state_S = torch.masked_select(state_S, mask).view(-1, mask.size(-1))  #(bs * select, dim)
         state_T = torch.masked_select(state_T, mask).view(-1, mask.size(-1))  # (bs * select, dim)