minor fix

sample.py

@@ -174,7 +174,7 @@ def sample(
             embedding_vocab_size=model_config["embedding_vocab_size"],
             learn_sigma=model_config["learn_sigma"],
             in_channels=data_config["data_dim"],
-        ).to(device).bfloat16
+        ).to(device)
 
         state_dict = find_model(ckpt_path)
         model.load_state_dict(state_dict)