mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Disable autocast for cpu to fix error. Remove unused precision arg. (#518)
When running on just cpu (intel), a call to torch.layer_norm would error with RuntimeError: expected scalar type BFloat16 but found Float Fix buggy device handling in model.py. Tested with scripts/dream.py --full_precision on just cpu on intel laptop. Works but slow at ~10s/it.
This commit is contained in:
@ -209,8 +209,7 @@ class AttnBlock(nn.Module):
|
||||
|
||||
h_ = torch.zeros_like(k, device=q.device)
|
||||
|
||||
device_type = 'mps' if q.device.type == 'mps' else 'cuda'
|
||||
if device_type == 'cuda':
|
||||
if q.device.type == 'cuda':
|
||||
stats = torch.cuda.memory_stats(q.device)
|
||||
mem_active = stats['active_bytes.all.current']
|
||||
mem_reserved = stats['reserved_bytes.all.current']
|
||||
@ -612,9 +611,8 @@ class Decoder(nn.Module):
|
||||
del h3
|
||||
|
||||
# prepare for up sampling
|
||||
device_type = 'mps' if h.device.type == 'mps' else 'cuda'
|
||||
gc.collect()
|
||||
if device_type == 'cuda':
|
||||
if h.device.type == 'cuda':
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# upsampling
|
||||
|
Reference in New Issue
Block a user