fix fp16

app.py

@@ -7,7 +7,6 @@ import torch
 
 import transformers
 import tokenizers
-from torch import autocast
 
 from sampling import CAIFSampler, TopKWithTemperatureSampler
 from generator import Generator
@@ -57,6 +56,10 @@ def inference(lm_model_name: str, cls_model_name: str, prompt: str, fp16: bool =
     generator.set_caif_sampler(caif_sampler)
     ordinary_sampler = TopKWithTemperatureSampler()
     generator.set_ordinary_sampler(ordinary_sampler)
+    if device == "cpu":
+        autocast = torch.cpu.amp.autocast
+    else:
+        autocast = torch.cuda.amp.autocast
     with autocast(fp16):
         sequences, tokens = generator.sample_sequences(
             num_samples=1,