Skip to content
Draft
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions RecommenderSystems/dlrm/dlrm_train_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -442,8 +442,7 @@ def __init__(
def build(self, labels, dense_fields, sparse_fields):
logits = self.module(dense_fields.to("cuda"), sparse_fields.to("cuda"))
loss = self.loss(logits, labels.to("cuda"))
reduce_loss = flow.mean(loss)
reduce_loss.backward()
loss.backward()
return reduce_loss.to("cpu")
Comment thread
ShawnXuan marked this conversation as resolved.
Outdated


Expand Down Expand Up @@ -481,7 +480,7 @@ def save_model(subdir):

opt = flow.optim.SGD(dlrm_module.parameters(), lr=args.learning_rate)
lr_scheduler = make_lr_scheduler(args, opt)
loss = flow.nn.BCEWithLogitsLoss(reduction="none").to("cuda")
loss = flow.nn.BCEWithLogitsLoss(reduction="mean").to("cuda")

if args.loss_scale_policy == "static":
grad_scaler = flow.amp.StaticGradScaler(1024)
Expand Down