Phillnet-Mini-Max / _source_bound.py
ayjays132's picture
Complete Phillnet Mini Text-Vision release v1.1.0
1e114b1 verified
Raw
History Blame Contribute Delete
696 Bytes
"""Internal helper for parameterless modules bound to the one source layer."""
from __future__ import annotations
import weakref
from torch import nn
from .source import DendroSourceLayer
class SourceBoundModule(nn.Module):
"""An ``nn.Module`` that references, but never registers, the source module."""
def __init__(self, source: DendroSourceLayer) -> None:
super().__init__()
self.__dict__["_source_ref"] = weakref.ref(source)
@property
def source(self) -> DendroSourceLayer:
source = self.__dict__["_source_ref"]()
if source is None:
raise RuntimeError("The owning DendroSourceLayer has been released")
return source