ethanlshen commited on
Commit
7e38c8b
·
verified ·
1 Parent(s): 5321ffc

Update superposed/llama/superposed_generation.py

Browse files
superposed/llama/superposed_generation.py CHANGED
@@ -68,7 +68,7 @@ class SuperposedLlama:
68
  )
69
  tokenizer = Tokenizer(model_path=tokenizer_path)
70
  model_args.vocab_size = tokenizer.n_words
71
- torch.set_default_dtype(torch.cuda.HalfTensor)
72
  # Set up superposed decoding
73
  model = SuperposedTransformer(model_args)
74
  model.load_state_dict(checkpoint, strict=False)
 
68
  )
69
  tokenizer = Tokenizer(model_path=tokenizer_path)
70
  model_args.vocab_size = tokenizer.n_words
71
+ torch.set_default_dtype(torch.float16)
72
  # Set up superposed decoding
73
  model = SuperposedTransformer(model_args)
74
  model.load_state_dict(checkpoint, strict=False)