Hello, If it’s a gradient vansihing problem, this can be solved using clipping gradient.
You can do this using by registering a simple backward hook.
clip_value = 0.5
for p in model.parameters():
    p.register_hook(lambda grad: torch.clamp(grad, -clip_value, clip_value))