Update block_baseline.py

block_baseline.py

@@ -23,7 +23,7 @@ def get_bins(vocab_size, block_size):
 
     return bin2words, words2bin
 
-def encode_block(model, enc, message, context, block_size, bin2words, words2bin, finish_sent=False, device='cuda'):
+def encode_block(model, enc, message, context, block_size, bin2words, words2bin, finish_sent=False, device='cpu'):
     length = len(message)
 
     context = torch.tensor(context[-1022:], device=device, dtype=torch.long)
@@ -96,7 +96,7 @@ def encode_block(model, enc, message, context, block_size, bin2words, words2bin,
 
     return output[len(context):].tolist(), avg_NLL, avg_KL, words_per_bit
 
-def decode_block(model, enc, text, context, block_size, bin2words, words2bin, device='cuda'):
+def decode_block(model, enc, text, context, block_size, bin2words, words2bin, device='cpu'):
     # inp is a list of token indices
     # context is a list of token indices
     inp = enc.encode(text)