Update arithmetic.py

arithmetic.py

@@ -3,7 +3,7 @@ import torch.nn.functional as F
 
 from utils import limit_past, kl, entropy, bits2int, int2bits, is_sent_finish, num_same_from_beg
 
-def encode_arithmetic(model, enc, message, context, finish_sent=False, device='cuda', temp=1.0, precision=16, topk=50000):
+def encode_arithmetic(model, enc, message, context, finish_sent=False, device='cpu', temp=1.0, precision=16, topk=50000):
 
     context = torch.tensor(context[-1022:], device=device, dtype=torch.long)
 
@@ -122,7 +122,7 @@ def encode_arithmetic(model, enc, message, context, finish_sent=False, device='c
 
     return output[len(context):].tolist(), avg_NLL, avg_KL, words_per_bit
 
-def decode_arithmetic(model, enc, text, context, device='cuda', temp=1.0, precision=16, topk=50000):
+def decode_arithmetic(model, enc, text, context, device='cpu', temp=1.0, precision=16, topk=50000):
     # inp is a list of token indices
     # context is a list of token indices
     inp = enc.encode(text)