apsys commited on
Commit
028eaf2
1 Parent(s): ab478b1

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1 -1
main.py CHANGED
@@ -61,7 +61,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
61
  model = AutoModelForCausalLM.from_pretrained(
62
  model_path,
63
  cache_dir=".",
64
- load_in_8bit=True if 'cuda' in device else False, # Enable loading in INT8
65
  device_map="auto" # Automatically map model to available devices
66
  )
67
 
 
61
  model = AutoModelForCausalLM.from_pretrained(
62
  model_path,
63
  cache_dir=".",
64
+ load_in_8bit=False, # Enable loading in INT8
65
  device_map="auto" # Automatically map model to available devices
66
  )
67