nyasukun commited on
Commit
47eb8cc
·
1 Parent(s): 11559c2

logging device

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -46,7 +46,10 @@ def generate_text_local(model_path, prompt, max_new_tokens=512, temperature=0.7,
46
  # Use the already initialized model
47
  if model_path in pipelines:
48
  model_pipeline = pipelines[model_path]
49
- logger.info(f"Running text generation with {model_path}")
 
 
 
50
 
51
  outputs = model_pipeline(
52
  prompt,
 
46
  # Use the already initialized model
47
  if model_path in pipelines:
48
  model_pipeline = pipelines[model_path]
49
+
50
+ # Log GPU usage information
51
+ device_info = next(model_pipeline.model.parameters()).device
52
+ logger.info(f"Running text generation with {model_path} on device: {device_info}")
53
 
54
  outputs = model_pipeline(
55
  prompt,