diff --git a/model.py b/model.py
index 404822e..1f9e7d1 100644
--- a/model.py
+++ b/model.py
@@ -471,7 +471,7 @@ class Tacotron2(nn.Module):
             output_lengths = batch
         text_padded = to_gpu(text_padded).long()
         input_lengths = to_gpu(input_lengths).long()
-        max_len = torch.max(input_lengths.data).item()
+        max_len = torch.max(input_lengths.data)
         mel_padded = to_gpu(mel_padded).float()
         gate_padded = to_gpu(gate_padded).float()
         output_lengths = to_gpu(output_lengths).long()