Merge pull request #12 from VE-FORBRYDERNE/breakmodel
Fix for when breakmodel layers is 0
This commit is contained in:
commit
03a4da0ade
|
@ -291,11 +291,12 @@ def new_forward(
|
||||||
gc.collect()
|
gc.collect()
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
for param1,param2 in zip(self.h[0].parameters(),self.extrastorage[0].parameters()):
|
if ram_blocks:
|
||||||
param1.data = param2.data.to(gpu_device, non_blocking=False).detach()
|
for param1,param2 in zip(self.h[0].parameters(),self.extrastorage[0].parameters()):
|
||||||
|
param1.data = param2.data.to(gpu_device, non_blocking=False).detach()
|
||||||
|
|
||||||
for param1,param2 in zip(self.h[ram_blocks-1].parameters(),self.extrastorage[ram_blocks-1].parameters()):
|
for param1,param2 in zip(self.h[ram_blocks-1].parameters(),self.extrastorage[ram_blocks-1].parameters()):
|
||||||
param1.data = param2.data.to(gpu_device, non_blocking=False).detach()
|
param1.data = param2.data.to(gpu_device, non_blocking=False).detach()
|
||||||
#END MODEL BREAK EDITS
|
#END MODEL BREAK EDITS
|
||||||
|
|
||||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||||
|
|
Loading…
Reference in New Issue