Solved by SaviourZ in post #4 
              
              
                This implement will have different output compared with tf.sequence_mask when the input is 2d-dimension. 
e.g When lengths is array([[1],[3], [2],[4]]). 
So I improved this implement as following: 
    def sequence_mask(self, lengths, maxlen=None, dtype=torch.bool):
        if maxlen is None:
      …