|
|
@ -115,7 +115,7 @@ public final class PNDMScheduler: Scheduler {
|
|
|
|
) {
|
|
|
|
) {
|
|
|
|
self.trainStepCount = trainStepCount
|
|
|
|
self.trainStepCount = trainStepCount
|
|
|
|
self.inferenceStepCount = stepCount
|
|
|
|
self.inferenceStepCount = stepCount
|
|
|
|
|
|
|
|
|
|
|
|
switch betaSchedule {
|
|
|
|
switch betaSchedule {
|
|
|
|
case .linear:
|
|
|
|
case .linear:
|
|
|
|
self.betas = linspace(betaStart, betaEnd, trainStepCount)
|
|
|
|
self.betas = linspace(betaStart, betaEnd, trainStepCount)
|
|
|
|