[Code scan] Respect CPU fallback before moving pretraining models to CUDA

[njzjz]

This issue is a result of a Codex global repository scan.

Summary

The non-DDP pretraining trainer calls self.model.cuda() unconditionally, then later checks whether CUDA is available and disables fp16. On CPU-only machines the trainer fails during construction before the fallback check can run.

Code references

unimol_tools/unimol_tools/pretrain/trainer.py

Lines 44 to 67 in 4596596

	# DDP setup
	if 'WORLD_SIZE' in os.environ and int(os.environ['WORLD_SIZE']) > 1:
	torch.cuda.set_device(self.local_rank)
	dist.init_process_group(backend='nccl')
	self.rank = dist.get_rank()
	self.model = self.model.to(self.local_rank)
	if self.fp16:
	self.model = self.model.half()
	if isinstance(self.loss_fn, nn.Module):
	self.loss_fn = self.loss_fn.to(self.local_rank).half()
	else:
	if isinstance(self.loss_fn, nn.Module):
	self.loss_fn = self.loss_fn.to(self.local_rank)
	self.model = DDP(self.model, device_ids=[self.local_rank])
	else:
	self.model = self.model.cuda()
	if self.fp16:
	self.model = self.model.half()
	if isinstance(self.loss_fn, nn.Module):
	self.loss_fn = self.loss_fn.cuda().half()
	else:
	if isinstance(self.loss_fn, nn.Module):
	self.loss_fn = self.loss_fn.cuda()

unimol_tools/unimol_tools/pretrain/trainer.py

Lines 164 to 174 in 4596596

	if self.fp16 and not torch.cuda.is_available():
	logger.warning("FP16 requested but CUDA is not available; disabling fp16.")
	self.fp16 = False
	self.scaler = (
	DynamicLossScaler(
	init_scale=getattr(config, "fp16_init_scale", 4),
	scale_window=getattr(config, "fp16_scale_window", 256),
	)
	if self.fp16
	else None
	)

Impact

Pretraining cannot run or be smoke-tested on CPU-only environments despite having later logic intended to disable fp16 when CUDA is unavailable.

Suggested fix

Choose self.device first, disable fp16 before moving modules, and use model.to(self.device) and loss_fn.to(self.device) instead of unconditional .cuda() calls.