MCPcopy Create free account
hub / github.com/csslc/PiSA-SR / _replace_module

Method _replace_module

src/models/autoencoder_kl.py:535–560  ·  view source on GitHub ↗
(self, parent, child_name, new_module, child)

Source from the content-addressed store, hash-verified

533 return self
534
535 def _replace_module(self, parent, child_name, new_module, child):
536 setattr(parent, child_name, new_module)
537 # It's not necessary to set requires_grad here, as that is handled by
538 # _mark_only_adapters_as_trainable
539
540 # child layer wraps the original module, unpack it
541 if hasattr(child, "base_layer"):
542 child = child.base_layer
543
544 if not hasattr(new_module, "base_layer"):
545 new_module.weight = child.weight
546 if hasattr(child, "bias"):
547 new_module.bias = child.bias
548
549 if getattr(child, "state", None) is not None:
550 if hasattr(new_module, "base_layer"):
551 new_module.base_layer.state = child.state
552 else:
553 new_module.state = child.state
554 new_module.to(child.weight.device)
555
556 # dispatch to correct device
557 for name, module in new_module.named_modules():
558 if ("lora_" in name) or ("ranknum" in name):
559 weight = child.qweight if hasattr(child, "qweight") else child.weight
560 module.to(weight.device)

Callers 1

Calls

no outgoing calls

Tested by

no test coverage detected