The treaty you're thinking of is the Treaty on the Non-Proliferation of Nuclear Weapons (also known as the NPT or NNPT). The core of this treaty for non-nuclear weapon states is Article II, which says
Each non-nuclear-weapon State Party to the Treaty undertakes not to receive the transfer from any transferor whatsoever of nuclear weapons or other nuclear explosive devices or of control over such weapons or explosive devices directly, or indirectly; not to manufacture or otherwise acquire nuclear weapons or other nuclear explosive devices; and not to seek or receive any assistance in the manufacture of nuclear weapons or other nuclear explosive devices.
As you can see, the limit is on "nuclear weapons or other nuclear explosive devices," not on all military applications of nuclear power. Another provision of the NPT reinforces this. Article III requires non-nuclear weapon states to enter into safeguards agreements with the IAEA to ensure that peaceful applications of nuclear energy aren't diverted to make nuclear weapons. In the guide to these agreements (IAEA INFCIRC/153), paragraph 14 covers use of nuclear material for military purposes. Countries are generally expected to notify the IAEA if they're doing this and give assurances that they won't use the material for weapons, but other than that it's perfectly fine to use the material for non-peaceful purposes. The country does not have to tell the IAEA any classified details about the use or get their approval for the military activity, just that X amount of Y material is being used for military purposes. Countries have used this provision when considering naval propulsion in the past.
Countries might also take the position that naval nuclear propulsion is a peaceful use (which would make them subject to safeguards, but might give advantages under other treaties). For instance, Brazil and Argentina entered into an agreement (INFCIRC/395) to exclusively use nuclear power for peaceful purposes, but in Article 3 declared that nuclear propulsion is a peaceful use. In their IAEA safeguards agreement (INFCIRC/435), Article 13 essentially restates paragraph 14 of INFCIRC/153 but changes some terms to be consistent with that position. "Material not subject to safeguards" becomes "material subject to special procedures" and "non-peaceful applications" becomes "nuclear propulsion," but it's otherwise basically the same. Whether or not nuclear propulsion is a peaceful use, this agreement (which the IAEA signed) shows that it's fully compatible with the NPT.
Because naval propulsion is exempt from detailed monitoring under safeguards agreements (all military use is, but naval propulsion is basically the only practical military use for nuclear energy other than weapons), it's been brought up as a potential loophole in the NPT. Naval propulsion also tends to involve much more highly-enriched uranium than civil power, so the concern is that a country could divert highly-enriched uranium from naval propulsion and use it to build weapons. However, the current rule is that naval propulsion is absolutely an acceptable use for nuclear power even by a non-nuclear weapon state.