_null_context in train.py is fragile and confusing #18

New issue

Open

opened 2026-05-09 19:22:08 +02:00 by sleepy · 0 comments

sleepy commented

2026-05-09 19:22:08 +02:00

Owner

Problem

train.py:185-194 uses a convoluted conditional for the autocast context:

with torch.amp.autocast(...) if self.device.type == "cuda" else \
     torch.no_grad() if False else _null_context():

The torch.no_grad() if False is dead code. _null_context.__enter__ returns None, which works but is confusing. If someone accidentally changes the condition, this could silently drop gradients.

Impact

Confusing code that's hard to maintain
Risk of accidentally wrapping the forward pass in torch.no_grad() if the condition is changed

Action needed

Replace with a clean, explicit context manager:

class _pass_through:
    def __enter__(self): return self
    def __exit__(self, *args): pass

if self.device.type == "cuda":
    ctx = torch.amp.autocast(device_type=self.device.type, dtype=self.amp_dtype)
else:
    ctx = _pass_through()
with ctx:
    ...

Files

tergent/train.py:185-194

## Problem `train.py:185-194` uses a convoluted conditional for the autocast context: ```python with torch.amp.autocast(...) if self.device.type == "cuda" else \ torch.no_grad() if False else _null_context(): ``` The `torch.no_grad() if False` is dead code. `_null_context.__enter__` returns `None`, which works but is confusing. If someone accidentally changes the condition, this could silently drop gradients. ## Impact - Confusing code that's hard to maintain - Risk of accidentally wrapping the forward pass in `torch.no_grad()` if the condition is changed ## Action needed Replace with a clean, explicit context manager: ```python class _pass_through: def __enter__(self): return self def __exit__(self, *args): pass if self.device.type == "cuda": ctx = torch.amp.autocast(device_type=self.device.type, dtype=self.amp_dtype) else: ctx = _pass_through() with ctx: ... ``` ## Files - `tergent/train.py:185-194`

No labels

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

sleepy/ternary#18

No description provided.

Rows
Columns